#!/usr/bin/env python3 from __future__ import annotations import argparse import os import sys from pathlib import Path import pymysql ROOT_DIR = Path(__file__).resolve().parent DEFAULT_SCHEMA_PATH = ROOT_DIR / "talkingq_shared_schema.sql" DEFAULT_ENV_PATH = ROOT_DIR.parent / "mini-program" / ".env" DEFAULT_TARGET_DB = "talkingq" DEFAULT_LEGACY_DB = "mini_program" TALKINGQ_URL_TABLES = [ "device_auth", "device_configs", "conversation_histories", "conversation_messages", "roles", "role_languages", "device_firmware_update", "system_config", ] MINI_PROGRAM_TABLES = [ "parents", "children", "parent_child_relations", "device_bindings", "device_bind_sessions", "device_bind_history", "cards", "device_settings", "im_conversations", "im_messages", "child_location_current", "child_location_history", ] def _env_first(*keys: str, default: str) -> str: for key in keys: value = os.getenv(key) if value: return value return default def load_env_file(path: Path) -> dict[str, str]: if not path.exists(): return {} env: dict[str, str] = {} for raw_line in path.read_text(encoding="utf-8").splitlines(): line = raw_line.strip() if not line or line.startswith("#") or "=" not in line: continue key, value = line.split("=", 1) key = key.strip() value = value.strip() if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: value = value[1:-1] env[key] = value return env def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Apply the shared Banban schema to the single source-of-truth database. " "Existing talkingq-url data in the target database is preserved; legacy " "mini_program data is ignored." ) ) parser.add_argument( "--env-file", default=str(DEFAULT_ENV_PATH), help="Optional env file used as fallback for DB_HOST/DB_PORT/DB_USER/DB_PASSWORD.", ) parser.add_argument("--host") parser.add_argument("--port", type=int) parser.add_argument("--user") parser.add_argument("--password") parser.add_argument("--target-db") parser.add_argument("--legacy-db") parser.add_argument("--schema", default=str(DEFAULT_SCHEMA_PATH)) parser.add_argument( "--drop-legacy-db", action="store_true", help="Drop the legacy mini_program database after the shared schema has been applied.", ) args = parser.parse_args() env_values = load_env_file(Path(args.env_file).expanduser().resolve()) args.host = args.host or os.getenv("MYSQL_HOST") or os.getenv("DB_HOST") or env_values.get("DB_HOST") or "127.0.0.1" args.port = ( args.port or int(os.getenv("MYSQL_PORT") or os.getenv("DB_PORT") or env_values.get("DB_PORT") or "3306") ) args.user = args.user or os.getenv("MYSQL_USER") or os.getenv("DB_USER") or env_values.get("DB_USER") or "root" args.password = ( args.password or os.getenv("MYSQL_PASSWORD") or os.getenv("DB_PASSWORD") or env_values.get("DB_PASSWORD") or "" ) args.target_db = args.target_db or DEFAULT_TARGET_DB args.legacy_db = args.legacy_db or DEFAULT_LEGACY_DB return args def split_sql_statements(sql: str) -> list[str]: statements: list[str] = [] buffer: list[str] = [] quote_char: str | None = None escaped = False for char in sql: buffer.append(char) if quote_char is not None: if escaped: escaped = False elif char == "\\" and quote_char in {"'", '"'}: escaped = True elif char == quote_char: quote_char = None continue if char in {"'", '"', "`"}: quote_char = char continue if char == ";": statement = "".join(buffer).strip() if statement: statements.append(statement) buffer = [] tail = "".join(buffer).strip() if tail: statements.append(tail) return statements def load_schema_sql(path: Path, *, target_db: str) -> str: raw_sql = path.read_text(encoding="utf-8-sig") sql_lines = [] for line in raw_sql.splitlines(): stripped = line.lstrip() if stripped.startswith("--"): continue sql_lines.append(line) sql = "\n".join(sql_lines) sql = sql.replace("CREATE DATABASE IF NOT EXISTS `talkingq`", f"CREATE DATABASE IF NOT EXISTS `{target_db}`") sql = sql.replace("USE `talkingq`;", f"USE `{target_db}`;") return sql def execute_statements(connection: pymysql.Connection, statements: list[str]) -> None: with connection.cursor() as cursor: for index, statement in enumerate(statements, start=1): try: cursor.execute(statement) except Exception as exc: # pragma: no cover - operational path snippet = " ".join(statement.split())[:200] raise RuntimeError(f"statement #{index} failed: {snippet}") from exc def fetch_table_names(connection: pymysql.Connection, *, database: str) -> set[str]: with connection.cursor() as cursor: cursor.execute( """ SELECT table_name FROM information_schema.tables WHERE table_schema = %s """, (database,), ) return {row[0] for row in cursor.fetchall()} def database_exists(connection: pymysql.Connection, *, database: str) -> bool: with connection.cursor() as cursor: cursor.execute("SHOW DATABASES LIKE %s", (database,)) return cursor.fetchone() is not None def drop_database(connection: pymysql.Connection, *, database: str) -> None: with connection.cursor() as cursor: cursor.execute(f"DROP DATABASE IF EXISTS `{database}`") def print_summary(connection: pymysql.Connection, *, target_db: str) -> None: tables = fetch_table_names(connection, database=target_db) expected_tables = set(TALKINGQ_URL_TABLES + MINI_PROGRAM_TABLES) missing_tables = sorted(expected_tables - tables) print(f"target database: {target_db}") print(f"talkingq-url tables present: {len(set(TALKINGQ_URL_TABLES) & tables)}/{len(TALKINGQ_URL_TABLES)}") print(f"mini-program tables present: {len(set(MINI_PROGRAM_TABLES) & tables)}/{len(MINI_PROGRAM_TABLES)}") if missing_tables: print("missing tables:") for table_name in missing_tables: print(f" - {table_name}") raise SystemExit(1) def main() -> int: args = parse_args() schema_path = Path(args.schema).resolve() if not schema_path.exists(): print(f"schema file not found: {schema_path}", file=sys.stderr) return 1 if args.drop_legacy_db and args.legacy_db == args.target_db: print("legacy database must be different from target database", file=sys.stderr) return 1 schema_sql = load_schema_sql(schema_path, target_db=args.target_db) statements = split_sql_statements(schema_sql) if not statements: print("schema file contains no executable statements", file=sys.stderr) return 1 connection = pymysql.connect( host=args.host, port=args.port, user=args.user, password=args.password, charset="utf8mb4", autocommit=True, ) try: print(f"applying shared schema to `{args.target_db}` using `{schema_path}`") print(f"db connection target: {args.user}@{args.host}:{args.port}") if args.env_file: print(f"env fallback file: {Path(args.env_file).expanduser().resolve()}") print("migration mode: talkingq-url data is preserved; legacy mini_program data is ignored") execute_statements(connection, statements) print_summary(connection, target_db=args.target_db) if args.drop_legacy_db: if database_exists(connection, database=args.legacy_db): drop_database(connection, database=args.legacy_db) print(f"dropped legacy database: {args.legacy_db}") else: print(f"legacy database not found, skip drop: {args.legacy_db}") else: print(f"legacy database left untouched: {args.legacy_db}") finally: connection.close() print("shared database migration completed") return 0 if __name__ == "__main__": raise SystemExit(main())