Files
tw/legacy/scripts/migrate_remaining_sqlite_to_mysql.py

110 lines
3.6 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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()