#!/usr/bin/env python3 from __future__ import annotations import argparse import hashlib import json import os import sqlite3 import sys from collections import defaultdict from collections.abc import Iterable, Sequence from pathlib import Path from typing import Any ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from backend.sql_catalog.database import ( # noqa: E402 MySQLConnection, connect_mysql, finding_run_idempotency_key, ) from backend.sql_catalog.tidb_schema import ( # noqa: E402 SCHEMA_VERSION, ensure_tidb_schema, ) DEFAULT_SOURCE = ROOT / "data" / "sql_catalog.sqlite3" ENTITY_TABLES = { "catalog_item": "catalog_items", "finding": "findings", "treatment": "treatments", "record": "records", } PRIMARY_KEYS: dict[str, tuple[str, ...]] = { "catalog_items": ("id",), "finding_runs": ("id",), "cluster_metadata": ("cluster_id",), "cluster_scopes": ("id",), "cluster_scope_members": ("id",), "user_preferences": ("owner_id",), "cluster_sql_health_owners": ("cluster_id", "owner_id"), "cluster_sql_review_policies": ("cluster_id",), "health_check_report_syncs": ("id",), "followed_scopes": ("id",), "user_scope_settings": ("owner_id",), "groups": ("id",), "group_members": ("id",), "symptom_definitions": ("code",), "finding_deletions": ("id",), "findings": ("id",), "finding_symptoms": ("id",), "treatments": ("id",), "records": ("id",), "relations": ("id",), "timeline_events": ("id",), } INSERT_ORDER = ( "symptom_definitions", "cluster_metadata", "cluster_scopes", "cluster_scope_members", "user_preferences", "user_scope_settings", "followed_scopes", "groups", "group_members", "cluster_sql_health_owners", "cluster_sql_review_policies", "finding_deletions", "catalog_items", "finding_runs", "health_check_report_syncs", "findings", "finding_symptoms", "treatments", "records", "relations", "timeline_events", ) def load_dotenv(path: Path) -> None: if not path.exists(): return 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) os.environ.setdefault(key.strip(), value.strip().strip("'").strip('"')) def placeholders(values: Sequence[Any]) -> str: return ",".join("?" for _ in values) def select_in( conn: sqlite3.Connection, table: str, column: str, values: Iterable[str], *, order_by: str | None = None, ) -> list[sqlite3.Row]: selected = sorted(set(values)) if not selected: return [] suffix = f" ORDER BY {order_by}" if order_by else "" return conn.execute( f"SELECT * FROM `{table}` WHERE `{column}` IN ({placeholders(selected)}){suffix}", selected, ).fetchall() def select_scope( conn: sqlite3.Connection, scope_id: str | None, scope_name: str ) -> sqlite3.Row: if scope_id: row = conn.execute( "SELECT * FROM cluster_scopes WHERE id = ?", (scope_id,) ).fetchone() else: rows = conn.execute( "SELECT * FROM cluster_scopes WHERE name = ? ORDER BY created_at", (scope_name,), ).fetchall() if len(rows) > 1: raise RuntimeError( f"Multiple scopes are named {scope_name!r}; pass --scope-id explicitly" ) row = rows[0] if rows else None if row is None: identity = scope_id or scope_name raise RuntimeError(f"Source scope not found: {identity}") return row def relation_closure( all_relations: Sequence[sqlite3.Row], entities: dict[str, set[str]] ) -> tuple[list[sqlite3.Row], dict[str, set[str]]]: selected_relation_ids: set[str] = set() changed = True while changed: changed = False for relation in all_relations: relation_id = str(relation["id"]) left_type = str(relation["from_type"]) left_id = str(relation["from_id"]) right_type = str(relation["to_type"]) right_id = str(relation["to_id"]) touches_selection = ( left_id in entities.get(left_type, set()) or right_id in entities.get(right_type, set()) ) if not touches_selection: continue if relation_id not in selected_relation_ids: selected_relation_ids.add(relation_id) changed = True if left_type in ENTITY_TABLES and left_id not in entities[left_type]: entities[left_type].add(left_id) changed = True if right_type in ENTITY_TABLES and right_id not in entities[right_type]: entities[right_type].add(right_id) changed = True return ( [row for row in all_relations if str(row["id"]) in selected_relation_ids], entities, ) def collect_scope_rows( source: sqlite3.Connection, scope_id: str | None, scope_name: str ) -> tuple[sqlite3.Row, dict[str, list[sqlite3.Row]]]: scope = select_scope(source, scope_id, scope_name) members = source.execute( "SELECT * FROM cluster_scope_members WHERE scope_id = ? " "ORDER BY cluster_name, cluster_id", (scope["id"],), ).fetchall() cluster_ids = {str(row["cluster_id"]) for row in members} if not cluster_ids: raise RuntimeError(f"Scope {scope['id']} has no member clusters") org_ids = {str(row["org_id"]) for row in members} catalogs = select_in(source, "catalog_items", "cluster_id", cluster_ids) findings = select_in(source, "findings", "cluster_id", cluster_ids) finding_runs = select_in(source, "finding_runs", "cluster_id", cluster_ids) report_syncs = select_in( source, "health_check_report_syncs", "cluster_id", cluster_ids ) entities: dict[str, set[str]] = defaultdict(set) entities["catalog_item"].update(str(row["id"]) for row in catalogs) entities["finding"].update(str(row["id"]) for row in findings) entities["treatment"] = set() entities["record"] = set() all_relations = source.execute("SELECT * FROM relations ORDER BY rowid").fetchall() relations, entities = relation_closure(all_relations, entities) # Relations can pull in a Treatment or Record attached to a selected # Catalog Item. Their full rows and timelines are part of the same dossier. treatments = select_in(source, "treatments", "id", entities["treatment"]) records = select_in(source, "records", "id", entities["record"]) finding_ids = {str(row["id"]) for row in findings} deletion_ids = { str(row["deletion_id"]) for row in findings if row["deletion_id"] is not None } timeline_rows = source.execute( "SELECT * FROM timeline_events ORDER BY rowid" ).fetchall() timeline = [ row for row in timeline_rows if str(row["target_id"]) in entities.get(str(row["target_type"]), set()) ] owner_ids = {str(scope["owner_id"])} preferences = source.execute( "SELECT * FROM user_preferences WHERE default_cluster_scope_id = ?", (scope["id"],), ).fetchall() owner_ids.update(str(row["owner_id"]) for row in preferences) user_scope_settings = select_in( source, "user_scope_settings", "owner_id", owner_ids ) followed_scopes = [ row for row in select_in(source, "followed_scopes", "owner_id", owner_ids) if ( row["scope_type"] == "cluster" and str(row["scope_id"]) in cluster_ids ) or ( row["scope_type"] == "organization" and str(row["org_id"]) in org_ids ) ] group_member_rows = source.execute( "SELECT * FROM group_members ORDER BY rowid" ).fetchall() group_ids = { str(row["group_id"]) for row in group_member_rows if row["member_type"] == "cluster" and str(row["member_id"]) in cluster_ids } groups = select_in(source, "groups", "id", group_ids) group_members = [ row for row in group_member_rows if str(row["group_id"]) in group_ids ] rows: dict[str, list[sqlite3.Row]] = { "symptom_definitions": source.execute( "SELECT * FROM symptom_definitions ORDER BY code" ).fetchall(), "cluster_metadata": select_in( source, "cluster_metadata", "cluster_id", cluster_ids ), "cluster_scopes": [scope], "cluster_scope_members": list(members), "user_preferences": preferences, "user_scope_settings": user_scope_settings, "followed_scopes": followed_scopes, "groups": groups, "group_members": group_members, "cluster_sql_health_owners": select_in( source, "cluster_sql_health_owners", "cluster_id", cluster_ids ), "cluster_sql_review_policies": select_in( source, "cluster_sql_review_policies", "cluster_id", cluster_ids ), "finding_deletions": select_in( source, "finding_deletions", "id", deletion_ids ), "catalog_items": catalogs, "finding_runs": finding_runs, "health_check_report_syncs": report_syncs, "findings": findings, "finding_symptoms": select_in( source, "finding_symptoms", "finding_id", finding_ids ), "treatments": treatments, "records": records, "relations": relations, "timeline_events": timeline, } return scope, rows def upsert_rows( target: MySQLConnection, table: str, rows: Sequence[sqlite3.Row], *, batch_size: int, ) -> None: if not rows: return columns = list(rows[0].keys()) if table == "finding_runs": columns.append("idempotency_key") quoted_columns = ", ".join(f"`{column}`" for column in columns) value_slots = ", ".join("?" for _ in columns) primary_keys = set(PRIMARY_KEYS[table]) mutable_columns = [column for column in columns if column not in primary_keys] if mutable_columns: updates = ", ".join( f"`{column}` = VALUES(`{column}`)" for column in mutable_columns ) conflict_clause = f" ON DUPLICATE KEY UPDATE {updates}" else: key = columns[0] conflict_clause = f" ON DUPLICATE KEY UPDATE `{key}` = VALUES(`{key}`)" sql = ( f"INSERT INTO `{table}` ({quoted_columns}) VALUES ({value_slots})" f"{conflict_clause}" ) values = [] for row in rows: row_values = [row[column] for column in rows[0].keys()] if table == "finding_runs": row_values.append( finding_run_idempotency_key( str(row["org_id"]), str(row["cluster_id"]), str(row["source"]), str(row["source_run_id"]), ) ) values.append(tuple(row_values)) for start in range(0, len(values), batch_size): target.executemany(sql, values[start : start + batch_size]) def content_digest(rows: Sequence[Any], columns: Sequence[str]) -> str: digest = hashlib.sha256() for row in rows: encoded = json.dumps( [row[column] for column in columns], ensure_ascii=False, separators=(",", ":"), ).encode("utf-8") digest.update(len(encoded).to_bytes(8, "big")) digest.update(encoded) return digest.hexdigest() def verify_rows( target: MySQLConnection, selected: dict[str, list[sqlite3.Row]] ) -> dict[str, int]: counts: dict[str, int] = {} for table in INSERT_ORDER: source_rows = selected[table] key_columns = PRIMARY_KEYS[table] columns = list(source_rows[0].keys()) if source_rows else list(key_columns) target_rows = target.execute( f"SELECT {', '.join(f'`{column}`' for column in columns)} " f"FROM `{table}`" ).fetchall() source_by_key = { tuple(row[column] for column in key_columns): row for row in source_rows } target_by_key = { tuple(row[column] for column in key_columns): row for row in target_rows } missing = set(source_by_key) - set(target_by_key) if missing: sample = sorted(missing, key=str)[:3] raise RuntimeError(f"Verification failed for {table}; missing keys: {sample}") ordered_keys = sorted(source_by_key, key=str) ordered_source = [source_by_key[key] for key in ordered_keys] ordered_target = [target_by_key[key] for key in ordered_keys] if content_digest(ordered_source, columns) != content_digest( ordered_target, columns ): raise RuntimeError(f"Content verification failed for {table}") counts[table] = len(source_rows) source_timeline_ids = [str(row["id"]) for row in selected["timeline_events"]] selected_timeline_ids = set(source_timeline_ids) target_timeline_ids = [ str(row["id"]) for row in target.execute( "SELECT id FROM timeline_events ORDER BY event_sequence" ).fetchall() if str(row["id"]) in selected_timeline_ids ] if source_timeline_ids != target_timeline_ids: raise RuntimeError("Timeline event order verification failed") source_run_hashes = { str(row["id"]): finding_run_idempotency_key( str(row["org_id"]), str(row["cluster_id"]), str(row["source"]), str(row["source_run_id"]), ) for row in selected["finding_runs"] } target_run_hashes = { str(row["id"]): bytes(row["idempotency_key"]) for row in target.execute( "SELECT id, idempotency_key FROM finding_runs" ).fetchall() if str(row["id"]) in source_run_hashes } if source_run_hashes != target_run_hashes: raise RuntimeError("Finding Run idempotency hash verification failed") broken_findings = target.execute( "SELECT COUNT(*) AS count FROM findings f " "LEFT JOIN catalog_items ci ON ci.id = f.catalog_item_id " "WHERE ci.id IS NULL" ).fetchone()["count"] broken_symptoms = target.execute( "SELECT COUNT(*) AS count FROM finding_symptoms fs " "LEFT JOIN findings f ON f.id = fs.finding_id WHERE f.id IS NULL" ).fetchone()["count"] broken_syncs = target.execute( "SELECT COUNT(*) AS count FROM health_check_report_syncs h " "LEFT JOIN finding_runs r ON r.id = h.finding_run_id " "WHERE h.finding_run_id IS NOT NULL AND r.id IS NULL" ).fetchone()["count"] if any((broken_findings, broken_symptoms, broken_syncs)): raise RuntimeError( "Referential verification failed: " f"findings={broken_findings}, symptoms={broken_symptoms}, syncs={broken_syncs}" ) return counts def migrate( source_path: Path, database_url: str, password_env: str, scope_id: str | None, scope_name: str, batch_size: int, dry_run: bool, verify_only: bool, ) -> dict[str, Any]: source = sqlite3.connect(f"file:{source_path}?mode=ro", uri=True) source.row_factory = sqlite3.Row try: scope, selected = collect_scope_rows(source, scope_id, scope_name) finally: source.close() summary: dict[str, Any] = { "scope": {"id": scope["id"], "name": scope["name"]}, "tables": {table: len(selected[table]) for table in INSERT_ORDER}, } if dry_run: return summary target = connect_mysql(database_url, password_env) try: if not verify_only: ensure_tidb_schema(target) for table in INSERT_ORDER: upsert_rows(target, table, selected[table], batch_size=batch_size) target.commit() summary["verified"] = verify_rows(target, selected) summary["content_verified"] = True summary["timeline_order_verified"] = True summary["finding_run_hashes_verified"] = True except Exception: target.rollback() raise finally: target.close() return summary def parser() -> argparse.ArgumentParser: result = argparse.ArgumentParser( description="Migrate one SQL Catalog Cluster Scope from SQLite to TiDB." ) result.add_argument("--source", type=Path, default=DEFAULT_SOURCE) result.add_argument("--scope-id") result.add_argument("--scope-name", required=True) result.add_argument( "--database-url-env", default="SQL_CATALOG_DATABASE_URL" ) result.add_argument("--password-env", default="DB_PASSWD") result.add_argument("--batch-size", type=int, default=100) result.add_argument("--dry-run", action="store_true") result.add_argument("--verify-only", action="store_true") return result def main() -> None: args = parser().parse_args() load_dotenv(ROOT / ".env") database_url = os.environ.get(args.database_url_env) if not database_url and not args.dry_run: raise SystemExit(f"Missing environment variable: {args.database_url_env}") if args.batch_size < 1 or args.batch_size > 1000: raise SystemExit("--batch-size must be between 1 and 1000") if args.dry_run and args.verify_only: raise SystemExit("--dry-run and --verify-only cannot be used together") summary = migrate( args.source.resolve(), database_url or "", args.password_env, args.scope_id, args.scope_name, args.batch_size, args.dry_run, args.verify_only, ) print(json.dumps(summary, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()