291 lines
8.7 KiB
Python
291 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
||
"""将本地 SQLite 数据迁移到当前 PostgreSQL（Docker）。
|
||
|
||
默认源库: /app/backend/neural_sync.db
|
||
默认目标库: 环境变量 DATABASE_URL
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
from datetime import datetime
|
||
from typing import Any
|
||
|
||
import asyncpg
|
||
|
||
TABLE_ORDER = [
|
||
"skill_tags",
|
||
"users",
|
||
"projects",
|
||
"feishu_items",
|
||
"signals",
|
||
"truth_nodes",
|
||
"tasks",
|
||
"project_signal_links",
|
||
"user_skill_bindings",
|
||
"user_workloads",
|
||
"qa_sessions",
|
||
"qa_records",
|
||
"qa_feedback",
|
||
]
|
||
|
||
|
||
def _parse_json(value: Any, default: Any) -> Any:
|
||
if value is None:
|
||
return default
|
||
if isinstance(value, (dict, list)):
|
||
return value
|
||
text = str(value).strip()
|
||
if not text:
|
||
return default
|
||
try:
|
||
return json.loads(text)
|
||
except Exception:
|
||
return default
|
||
|
||
|
||
def _parse_ts(value: Any) -> Any:
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, datetime):
|
||
return value
|
||
text = str(value).strip()
|
||
if not text:
|
||
return None
|
||
text = text.replace("T", " ").replace("Z", "")
|
||
try:
|
||
return datetime.fromisoformat(text)
|
||
except Exception:
|
||
return text
|
||
|
||
|
||
def _parse_int_array(value: Any) -> list[int]:
|
||
if value is None:
|
||
return []
|
||
if isinstance(value, list):
|
||
raw = value
|
||
else:
|
||
raw = _parse_json(value, default=[])
|
||
if not isinstance(raw, list):
|
||
raw = []
|
||
result: list[int] = []
|
||
for item in raw:
|
||
try:
|
||
num = int(item)
|
||
result.append(num)
|
||
except Exception:
|
||
continue
|
||
return result
|
||
|
||
|
||
def _transform(table: str, row: dict[str, Any], pg_types: dict[str, str]) -> dict[str, Any]:
|
||
data = dict(row)
|
||
|
||
if table == "users":
|
||
data["assigned_domains"] = _parse_json(data.get("assigned_domains"), default=[])
|
||
|
||
if table == "truth_nodes":
|
||
content = data.get("content")
|
||
parsed = _parse_json(content, default=None)
|
||
if parsed is None:
|
||
parsed = {"raw": str(content or "")}
|
||
data["content"] = parsed
|
||
|
||
if table == "tasks":
|
||
data["required_skill_tag_ids"] = _parse_int_array(data.get("required_skill_tag_ids"))
|
||
data["ai_recommended_assignees"] = _parse_json(data.get("ai_recommended_assignees"), default=[])
|
||
|
||
if table == "qa_records":
|
||
data["sources"] = _parse_json(data.get("sources"), default=[])
|
||
data["domains_used"] = _parse_json(data.get("domains_used"), default=[])
|
||
|
||
if table == "qa_feedback":
|
||
data["error_types"] = _parse_json(data.get("error_types"), default=[])
|
||
|
||
# 通用类型修正
|
||
for col, typ in pg_types.items():
|
||
if col not in data:
|
||
continue
|
||
if typ in {"json", "jsonb"}:
|
||
value = data[col]
|
||
if value is None:
|
||
continue
|
||
if isinstance(value, str):
|
||
data[col] = value
|
||
else:
|
||
data[col] = json.dumps(value, ensure_ascii=False)
|
||
continue
|
||
if typ == "boolean":
|
||
value = data[col]
|
||
if isinstance(value, bool):
|
||
continue
|
||
if value is None:
|
||
data[col] = False
|
||
else:
|
||
data[col] = str(value).strip().lower() in {"1", "true", "t", "yes", "y"}
|
||
elif typ == "timestamp without time zone":
|
||
data[col] = _parse_ts(data.get(col))
|
||
|
||
return data
|
||
|
||
|
||
async def _pg_table_columns(conn: asyncpg.Connection, table: str) -> list[str]:
|
||
rows = await conn.fetch(
|
||
"""
|
||
SELECT column_name
|
||
FROM information_schema.columns
|
||
WHERE table_schema='public' AND table_name=$1
|
||
ORDER BY ordinal_position
|
||
""",
|
||
table,
|
||
)
|
||
return [str(r["column_name"]) for r in rows]
|
||
|
||
|
||
async def _pg_table_types(conn: asyncpg.Connection, table: str) -> dict[str, str]:
|
||
rows = await conn.fetch(
|
||
"""
|
||
SELECT column_name, data_type
|
||
FROM information_schema.columns
|
||
WHERE table_schema='public' AND table_name=$1
|
||
""",
|
||
table,
|
||
)
|
||
return {str(r["column_name"]): str(r["data_type"]) for r in rows}
|
||
|
||
|
||
async def _pg_existing_tables(conn: asyncpg.Connection) -> set[str]:
|
||
rows = await conn.fetch(
|
||
"""
|
||
SELECT tablename
|
||
FROM pg_tables
|
||
WHERE schemaname='public'
|
||
"""
|
||
)
|
||
return {str(r["tablename"]) for r in rows}
|
||
|
||
|
||
def _sqlite_columns(sqlite_conn: sqlite3.Connection, table: str) -> list[str]:
|
||
rows = sqlite_conn.execute(f"PRAGMA table_info({table})").fetchall()
|
||
return [str(r[1]) for r in rows]
|
||
|
||
|
||
def _sqlite_rows(sqlite_conn: sqlite3.Connection, table: str) -> list[dict[str, Any]]:
|
||
sqlite_conn.row_factory = sqlite3.Row
|
||
rows = sqlite_conn.execute(f"SELECT * FROM {table} ORDER BY id ASC").fetchall()
|
||
return [dict(r) for r in rows]
|
||
|
||
|
||
async def migrate(sqlite_path: str, postgres_url: str, *, mode: str = "merge") -> None:
|
||
sqlite_conn = sqlite3.connect(sqlite_path)
|
||
sqlite_conn.row_factory = sqlite3.Row
|
||
|
||
pg_url = postgres_url.replace("postgresql+asyncpg://", "postgresql://")
|
||
pg = await asyncpg.connect(pg_url)
|
||
|
||
try:
|
||
pg_tables = await _pg_existing_tables(pg)
|
||
sqlite_tables = {
|
||
r[0]
|
||
for r in sqlite_conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
|
||
}
|
||
|
||
tables = [t for t in TABLE_ORDER if t in pg_tables and t in sqlite_tables]
|
||
if not tables:
|
||
raise RuntimeError("无可迁移表，请检查库路径与目标库")
|
||
|
||
print(f"[INFO] migrate tables: {tables}")
|
||
print(f"[INFO] mode: {mode}")
|
||
|
||
if mode == "replace":
|
||
# 清空目标核心表（含级联）
|
||
await pg.execute("TRUNCATE TABLE " + ", ".join(tables) + " RESTART IDENTITY CASCADE")
|
||
|
||
for table in tables:
|
||
pg_cols = await _pg_table_columns(pg, table)
|
||
pg_types = await _pg_table_types(pg, table)
|
||
sqlite_cols = _sqlite_columns(sqlite_conn, table)
|
||
common_cols = [c for c in pg_cols if c in sqlite_cols]
|
||
if not common_cols:
|
||
print(f"[WARN] skip {table}: no common cols")
|
||
continue
|
||
|
||
rows = _sqlite_rows(sqlite_conn, table)
|
||
inserted = 0
|
||
for row in rows:
|
||
row = _transform(table, row, pg_types)
|
||
values = [row.get(c) for c in common_cols]
|
||
placeholders = ", ".join(f"${i+1}" for i in range(len(common_cols)))
|
||
cols_sql = ", ".join(common_cols)
|
||
if mode == "merge" and "id" in common_cols and "id" in pg_cols:
|
||
sql = (
|
||
f"INSERT INTO {table} ({cols_sql}) VALUES ({placeholders}) "
|
||
"ON CONFLICT (id) DO NOTHING"
|
||
)
|
||
else:
|
||
sql = f"INSERT INTO {table} ({cols_sql}) VALUES ({placeholders})"
|
||
try:
|
||
await pg.execute(sql, *values)
|
||
inserted += 1
|
||
except Exception as exc:
|
||
print(f"[WARN] {table} row id={row.get('id')} failed: {exc}")
|
||
|
||
if "id" in common_cols:
|
||
await pg.execute(
|
||
f"SELECT setval(pg_get_serial_sequence('{table}','id'), COALESCE((SELECT MAX(id) FROM {table}), 1), true)"
|
||
)
|
||
|
||
print(f"[OK] {table}: {inserted}/{len(rows)}")
|
||
|
||
# 摘要
|
||
summary_tables = [
|
||
"users",
|
||
"projects",
|
||
"signals",
|
||
"truth_nodes",
|
||
"tasks",
|
||
"qa_sessions",
|
||
"qa_records",
|
||
"qa_feedback",
|
||
]
|
||
print("[SUMMARY]")
|
||
for table in summary_tables:
|
||
if table not in tables:
|
||
continue
|
||
count = await pg.fetchval(f"SELECT count(*) FROM {table}")
|
||
print(f" - {table}: {int(count or 0)}")
|
||
finally:
|
||
await pg.close()
|
||
sqlite_conn.close()
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--sqlite", default="/app/backend/neural_sync.db")
|
||
parser.add_argument("--postgres", default=os.getenv("DATABASE_URL", ""))
|
||
parser.add_argument("--mode", choices=["merge", "replace"], default="merge")
|
||
args = parser.parse_args()
|
||
|
||
if not args.postgres:
|
||
print("[ERROR] DATABASE_URL 为空，请使用 --postgres 指定")
|
||
return 2
|
||
|
||
if not os.path.exists(args.sqlite):
|
||
print(f"[ERROR] SQLite 文件不存在: {args.sqlite}")
|
||
return 2
|
||
|
||
try:
|
||
asyncio.run(migrate(args.sqlite, args.postgres, mode=args.mode))
|
||
return 0
|
||
except Exception as exc:
|
||
print(f"[ERROR] migration failed: {exc}")
|
||
return 2
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|