assistant/scripts/migrate_mysql_to_pg.py

115 lines
5.2 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.

"""一次性迁移:把 2026-09-16/17 期间 MySQL 里的数据搬到 PostgreSQL。
用法(两个库都在跑:docker compose --profile legacy up -d mysql;docker compose up -d postgres):
.venv\\Scripts\\python.exe scripts\\migrate_mysql_to_pg.py [--mysql-port 13307]
幂等:目标端用 ON CONFLICT DO NOTHING;knowledge 不迁移(启动时由 YAML 重新灌入)。
"""
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from dotenv import load_dotenv # noqa: E402
load_dotenv(ROOT / ".env")
import pymysql # noqa: E402
from pymysql.cursors import DictCursor # noqa: E402
from psycopg.types.json import Jsonb # noqa: E402
from backend.db import database as db # noqa: E402
def _json(value):
if value is None:
return None
if isinstance(value, (dict, list)):
return Jsonb(value)
try:
return Jsonb(json.loads(value))
except (TypeError, ValueError):
return None
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--mysql-host", default="127.0.0.1")
ap.add_argument("--mysql-port", type=int, default=13307)
ap.add_argument("--mysql-user", default="assistant")
ap.add_argument("--mysql-password", default="assistant_dev_123")
ap.add_argument("--mysql-db", default="school_assistant")
args = ap.parse_args()
src = pymysql.connect(
host=args.mysql_host, port=args.mysql_port, user=args.mysql_user, password=args.mysql_password,
database=args.mysql_db, charset="utf8mb4", cursorclass=DictCursor,
)
db.init_db()
counts: dict[str, int] = {}
def rows(sql: str) -> list[dict]:
with src.cursor() as cur:
cur.execute(sql)
return list(cur.fetchall())
with db.get_conn() as pg:
for r in rows("SELECT * FROM settings"):
counts["settings"] = counts.get("settings", 0) + db._execute(
pg, 'INSERT INTO settings ("key", "value") VALUES (%s, %s) ON CONFLICT DO NOTHING', (r["key"], r["value"]))
for r in rows("SELECT * FROM tasks"):
counts["tasks"] = counts.get("tasks", 0) + db._execute(
pg,
"""INSERT INTO tasks (id, school_id, student_id, knowledge_id, status, current_step, inputs, artifacts,
skill_version, browser_profile_path, pause_reason, next_check_at, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING""",
(r["id"], r.get("school_id") or "default", r["student_id"], r["knowledge_id"], r["status"], r["current_step"],
_json(r["inputs"]) or Jsonb({}), _json(r["artifacts"]) or Jsonb({}), r["skill_version"],
r["browser_profile_path"], r["pause_reason"], r["next_check_at"], r["created_at"], r["updated_at"]),
)
for r in rows("SELECT * FROM run_logs ORDER BY id"):
counts["run_logs"] = counts.get("run_logs", 0) + db._execute(
pg,
"""INSERT INTO run_logs (id, task_id, step_id, agent_history, screenshots, verdict, started_at, finished_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING""",
(r["id"], r["task_id"], r["step_id"], _json(r["agent_history"]), _json(r["screenshots"]) or Jsonb([]),
r["verdict"], r["started_at"], r["finished_at"]),
)
for r in rows("SELECT * FROM system_accounts ORDER BY id"):
counts["system_accounts"] = counts.get("system_accounts", 0) + db._execute(
pg,
"""INSERT INTO system_accounts (id, school_id, student_id, system_id, role, username, password_enc, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING""",
(r["id"], r.get("school_id") or "default", r["student_id"], r["system_id"], r.get("role") or "",
r["username"], r["password_enc"], r["created_at"], r["updated_at"]),
)
for r in rows("SELECT * FROM learned_entries ORDER BY id"):
counts["learned_entries"] = counts.get("learned_entries", 0) + db._execute(
pg,
"""INSERT INTO learned_entries (id, school_id, knowledge_id, business_id, "key", "value", status, category,
observations, source_task_id, created_at, updated_at, verified_at, verified_by)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING""",
(r["id"], r.get("school_id") or "default", r["knowledge_id"], r["business_id"], r["key"], r["value"],
r["status"], r.get("category") or "note", r.get("observations") or 1, r["source_task_id"],
r["created_at"], r["updated_at"], r["verified_at"], r.get("verified_by")),
)
for table in ("run_logs", "system_accounts", "learned_entries"):
db._execute(pg, f"SELECT setval(pg_get_serial_sequence('{table}', 'id'), COALESCE((SELECT MAX(id) FROM {table}), 0) + 1, false)")
src.close()
print("migrated rows:", counts)
return 0
if __name__ == "__main__":
raise SystemExit(main())