110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
迁移其余 SQLite 业务库到 MySQL(保留主键):
|
||
- deal_outcome_db/outcomes.db -> deal_outcomes
|
||
- designer_roster_db/roster.db -> designers/designer_shops/designer_online/round_robin
|
||
- image_tasks.db -> image_tasks/requirement_history
|
||
- task_db/tasks.db -> tasks/task_logs
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import sqlite3
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
import pymysql
|
||
|
||
|
||
MAPPINGS = [
|
||
{"sqlite": Path("db/deal_outcome_db/outcomes.db"), "tables": ["deal_outcomes"]},
|
||
{"sqlite": Path("db/designer_roster_db/roster.db"), "tables": ["designers", "designer_shops", "designer_online", "round_robin"]},
|
||
{"sqlite": Path("db/image_tasks.db"), "tables": ["image_tasks", "task_requirement_changes"]},
|
||
{"sqlite": Path("db/task_db/tasks.db"), "tables": ["tasks"]},
|
||
]
|
||
|
||
|
||
def mysql_conn(host: str, port: int, user: str, password: str, database: str):
|
||
return pymysql.connect(
|
||
host=host,
|
||
port=port,
|
||
user=user,
|
||
password=password,
|
||
database=database,
|
||
charset="utf8mb4",
|
||
autocommit=False,
|
||
cursorclass=pymysql.cursors.DictCursor,
|
||
)
|
||
|
||
|
||
def sqlite_table_exists(conn: sqlite3.Connection, table: str) -> bool:
|
||
row = conn.execute(
|
||
"SELECT name FROM sqlite_master WHERE type='table' AND name=?",
|
||
(table,),
|
||
).fetchone()
|
||
return row is not None
|
||
|
||
|
||
def sqlite_fetch_all(conn: sqlite3.Connection, table: str) -> List[sqlite3.Row]:
|
||
conn.row_factory = sqlite3.Row
|
||
return conn.execute(f"SELECT * FROM {table}").fetchall()
|
||
|
||
|
||
def migrate_table(mysql, rows: List[sqlite3.Row], table: str, truncate_target: bool) -> int:
|
||
if not rows:
|
||
return 0
|
||
cols = list(rows[0].keys())
|
||
col_sql = ", ".join(cols)
|
||
val_sql = ", ".join(["%s"] * len(cols))
|
||
sql = f"REPLACE INTO {table} ({col_sql}) VALUES ({val_sql})"
|
||
if truncate_target:
|
||
with mysql.cursor() as cur:
|
||
try:
|
||
cur.execute(f"TRUNCATE TABLE {table}")
|
||
except Exception:
|
||
try:
|
||
cur.execute(f"DELETE FROM {table}")
|
||
except Exception:
|
||
return 0
|
||
values = [tuple(r[c] for c in cols) for r in rows]
|
||
with mysql.cursor() as cur:
|
||
cur.executemany(sql, values)
|
||
mysql.commit()
|
||
return len(values)
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser(description="迁移剩余 SQLite 业务库到 MySQL")
|
||
p.add_argument("--host", required=True)
|
||
p.add_argument("--port", type=int, default=3306)
|
||
p.add_argument("--user", required=True)
|
||
p.add_argument("--password", required=True)
|
||
p.add_argument("--database", required=True)
|
||
p.add_argument("--truncate-target", action="store_true")
|
||
args = p.parse_args()
|
||
|
||
total = 0
|
||
with mysql_conn(args.host, args.port, args.user, args.password, args.database) as mconn:
|
||
for item in MAPPINGS:
|
||
sp = item["sqlite"]
|
||
if not sp.exists():
|
||
continue
|
||
sconn = sqlite3.connect(str(sp))
|
||
try:
|
||
for table in item["tables"]:
|
||
if not sqlite_table_exists(sconn, table):
|
||
continue
|
||
rows = sqlite_fetch_all(sconn, table)
|
||
n = migrate_table(mconn, rows, table, truncate_target=bool(args.truncate_target))
|
||
total += n
|
||
print(f"[MIGRATE] {sp}::{table} -> {n}")
|
||
finally:
|
||
sconn.close()
|
||
print(f"[DONE] migrated total rows: {total}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|