#!/usr/bin/env python3 from __future__ import annotations import argparse import concurrent.futures from dataclasses import dataclass from datetime import datetime, timezone from decimal import Decimal import hashlib import json import os from pathlib import Path import re import subprocess import tempfile import ssl import time from typing import Any DEFAULT_BULK_OBJECTS = ( "Account", "Contact", "Lead", "Campaign", "CampaignMemberStatus", "Pricebook2", "Product2", "PricebookEntry", "Profile", "UserRole", "PermissionSet", "User", "PermissionSetAssignment", "ProcessDefinition", "ProcessNode", "ProcessAction", "Quote", "ProcessInstance", "ProcessInstanceNode", "Report", "ReportFolder", "Dashboard", "_System_EmailTemplate", "_System_ConnectedApp", "HubSpotOwnerBinding__c", "Partner_Region_Owner_Rule__c", "Countries_Region__c", "Country_Dictionary__c", "Industry_Vertical__c", ) AUDIT_DATETIME_COLUMNS = ("CreatedDate", "LastModifiedDate", "SystemModstamp") AUDIT_USER_COLUMNS = ("CreatedById", "LastModifiedById", "OwnerId") NEVER_UPDATE_COLUMNS = {"Id", "CreatedDate", "CreatedById"} SENSITIVE_STATE_FIELDS = {"password"} TEMPLATE_RE = re.compile(r"\{\{\s*([^}]+?)\s*\}\}") PASSWORD_HASH_CACHE: dict[str, str] = {} PYMYSQL: Any | None = None class BulkSeedError(RuntimeError): pass class MissingReference(BulkSeedError): def __init__(self, key: str, expression: str): super().__init__(f"missing record reference {key!r} in {expression!r}") self.key = key self.expression = expression @dataclass(frozen=True) class DbConfig: host: str port: int user: str password: str database: str tls: bool def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Apply table-backed seed records through bulk TiDB writes.") seed_source = parser.add_mutually_exclusive_group(required=True) seed_source.add_argument("--seed-json", help="JSON rendering of the seed YAML.") seed_source.add_argument("--seed-file", help="Seed YAML file to render and apply.") parser.add_argument( "--state-file", required=True, help="Output runtime state JSON for evidence and adapter normalization.", ) parser.add_argument("--summary-file", required=True, help="Output bulk apply summary JSON.") parser.add_argument("--objects", help="Comma-separated bulk object allowlist.") parser.add_argument("--record-keys", help="Comma-separated baseline record key allowlist.") parser.add_argument( "--selected-only", action="store_true", help="Drop seed records outside --objects before planning, so coverage applies only to the selected objects.", ) parser.add_argument("--dry-run", action="store_true", help="Plan only; do not connect to TiDB or write rows.") parser.add_argument("--max-workers", type=int, default=4, help="Maximum parallel object writes per dependency layer.") parser.add_argument("--db-host", default=os.environ.get("TIDB_HOST", "")) parser.add_argument("--db-port", type=int, default=int(os.environ.get("TIDB_PORT", "4000"))) parser.add_argument("--db-user", default=os.environ.get("TIDB_USER", "")) parser.add_argument("--db-password", default=os.environ.get("TIDB_PASSWORD", "")) parser.add_argument("--db-database", default=os.environ.get("TIDB_DATABASE", "")) parser.add_argument("--db-tls", default=os.environ.get("TIDB_TLS", "true")) return parser.parse_args() def require_pymysql(): global PYMYSQL if PYMYSQL is None: try: import pymysql as pymysql_module except ImportError as exc: raise BulkSeedError("PyMySQL is required outside --dry-run") from exc PYMYSQL = pymysql_module return PYMYSQL def load_seed_payload(args: argparse.Namespace) -> dict[str, Any]: if args.seed_json: payload = json.loads(Path(args.seed_json).read_text(encoding="utf-8")) else: try: import yaml except ImportError as exc: raise BulkSeedError("PyYAML is required when using --seed-file") from exc payload = yaml.safe_load(Path(args.seed_file).read_text(encoding="utf-8")) if not isinstance(payload, dict): raise BulkSeedError("seed payload must be a YAML/JSON object") return payload def seed_payload_for_selected_objects(seed: dict[str, Any], selected_objects: set[str]) -> dict[str, Any]: filtered = dict(seed) filtered["baseline_records"] = [ record for record in seed.get("baseline_records") or [] if str(record.get("object") or "").strip() in selected_objects ] return filtered def seed_payload_for_selected_record_keys( seed: dict[str, Any], selected_record_keys: set[str] ) -> dict[str, Any]: available_record_keys = { str(record.get("key") or "").strip() for record in seed.get("baseline_records") or [] } unknown_record_keys = selected_record_keys - available_record_keys if unknown_record_keys: raise BulkSeedError( "unknown --record-keys: " + ", ".join(sorted(unknown_record_keys)) ) filtered = dict(seed) filtered["baseline_records"] = [ record for record in seed.get("baseline_records") or [] if str(record.get("key") or "").strip() in selected_record_keys ] return filtered def parse_list(raw: str | None, default: tuple[str, ...]) -> list[str]: if raw is None or raw.strip() == "": return list(default) return [item.strip() for item in raw.split(",") if item.strip()] def env_bool(name: str, default: bool = False) -> bool: raw = os.environ.get(name) if raw is None or raw.strip() == "": return default return raw.strip().lower() in {"1", "true", "yes", "on"} def tidb_tls_context() -> ssl.SSLContext: context = ssl.create_default_context() context.check_hostname = False context.verify_mode = ssl.CERT_NONE return context def db_config_from_args(args: argparse.Namespace) -> DbConfig: missing = [ name for name, value in [ ("TIDB_HOST", args.db_host), ("TIDB_USER", args.db_user), ("TIDB_DATABASE", args.db_database), ] if not str(value or "").strip() ] if missing: raise BulkSeedError(f"missing database settings: {', '.join(missing)}") return DbConfig( host=args.db_host, port=args.db_port, user=args.db_user, password=args.db_password, database=args.db_database, tls=str(args.db_tls).strip().lower() in {"", "1", "true", "yes", "on"}, ) def connect(db: DbConfig): pymysql = require_pymysql() kwargs: dict[str, Any] = { "host": db.host, "port": db.port, "user": db.user, "password": db.password, "database": db.database, "charset": "utf8mb4", "autocommit": False, "connect_timeout": 10, "read_timeout": 60, "write_timeout": 60, } if db.tls: kwargs["ssl"] = tidb_tls_context() return pymysql.connect(**kwargs) def stable_uuid(namespace: str, key: str) -> str: digest = hashlib.sha1(f"{namespace}:{key}".encode("utf-8")).hexdigest()[:32] chars = list(digest) chars[12] = "5" chars[16] = "8" return "-".join( [ "".join(chars[0:8]), "".join(chars[8:12]), "".join(chars[12:16]), "".join(chars[16:20]), "".join(chars[20:32]), ] ) def resolve_path(source: Any, path_expr: str) -> Any: current = source normalized = re.sub(r"\[(\d+)\]", r".\1", path_expr) for part in [part for part in normalized.split(".") if part]: if current is None: return None if isinstance(current, dict): current = current.get(part) continue if isinstance(current, list) and part.isdigit(): idx = int(part) current = current[idx] if idx < len(current) else None continue return None return current def resolve_expr(expr: str, context: dict[str, Any]) -> Any: expr = expr.strip() if expr == "run_id": return context.get("run_id") if expr.startswith("env."): key = expr[len("env.") :] value = os.environ.get(key) if value is None: raise BulkSeedError(f"missing environment variable for seed interpolation: {key}") return value if expr.startswith("inputs.") or expr.startswith("input."): path_expr = re.sub(r"^inputs?\.", "", expr) value = resolve_path(context.get("inputs", {}), path_expr) if value is None: raise BulkSeedError(f"missing input reference: {path_expr}") return value if expr.startswith("lookup.") or expr.startswith("lookups."): path_expr = re.sub(r"^lookups?\.", "", expr) value = resolve_path(context.get("lookups", {}), path_expr) if value is None: raise BulkSeedError(f"missing lookup reference: {path_expr}") return value if expr.startswith("record.") or expr.startswith("records."): path_expr = re.sub(r"^records?\.", "", expr) record_key = path_expr.split(".", 1)[0] value = resolve_path(context.get("records", {}), path_expr) if value is None: raise MissingReference(record_key, expr) return value return None def materialize_value(value: Any, context: dict[str, Any]) -> Any: if isinstance(value, str): full_match = TEMPLATE_RE.fullmatch(value.strip()) if full_match: resolved = resolve_expr(full_match.group(1), context) return value if resolved is None else resolved def replace(match: re.Match[str]) -> str: resolved = resolve_expr(match.group(1), context) return match.group(0) if resolved is None else str(resolved) return TEMPLATE_RE.sub(replace, value) if isinstance(value, list): return [materialize_value(item, context) for item in value] if isinstance(value, dict): return {key: materialize_value(nested, context) for key, nested in value.items()} return value def build_inputs(seed: dict[str, Any], run_id: str) -> dict[str, Any]: inputs: dict[str, Any] = {} context = { "seed": seed, "run_id": run_id, "inputs": inputs, "lookups": {}, "records": {}, } for key, raw_value in (seed.get("inputs") or {}).items(): inputs[key] = materialize_value(raw_value, context) inputs["validation_mode"] = "baseline" return inputs def stable_object(value: Any) -> Any: if isinstance(value, list): return [stable_object(item) for item in value] if isinstance(value, dict): return {key: stable_object(value[key]) for key in sorted(value)} return value def match_signature(object_api_name: str, match_fields: dict[str, Any]) -> str: return json.dumps({"object": object_api_name, "match": stable_object(match_fields)}, sort_keys=True, separators=(",", ":")) def state_for_row(row: dict[str, Any]) -> dict[str, Any]: values = { key: value for key, value in {**row["match"], **row["fields"]}.items() if key.lower() not in SENSITIVE_STATE_FIELDS } return { "id": row["id"], "object": row["object"], **values, } def row_by_key(records: list[dict[str, Any]]) -> dict[str, str]: result: dict[str, str] = {} for record in records: key = str(record.get("key") or "").strip() if key: result[key] = str(record.get("object") or "").strip() return result def records_by_object(seed: dict[str, Any], selected_objects: set[str]) -> dict[str, list[dict[str, Any]]]: grouped: dict[str, list[dict[str, Any]]] = {} for record in seed.get("baseline_records") or []: object_api_name = str(record.get("object") or "").strip() if object_api_name in selected_objects: grouped.setdefault(object_api_name, []).append(record) return grouped def unselected_records(seed: dict[str, Any], selected_objects: set[str]) -> list[dict[str, Any]]: return [ record for record in seed.get("baseline_records") or [] if str(record.get("object") or "").strip() not in selected_objects ] def unsupported_unselected_objects(seed: dict[str, Any], selected_objects: set[str]) -> list[dict[str, Any]]: grouped: dict[str, int] = {} for record in unselected_records(seed, selected_objects): object_api_name = str(record.get("object") or "").strip() grouped[object_api_name] = grouped.get(object_api_name, 0) + 1 return [ { "object": object_api_name, "reason": "not in bulk allowlist", "records": count, } for object_api_name, count in sorted(grouped.items()) ] def plan_one_object( object_api_name: str, object_records: list[dict[str, Any]], context: dict[str, Any], ) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]], int, int, int]: rows_by_signature: dict[str, dict[str, Any]] = {} ordered_signatures: list[str] = [] conflicting_signatures: set[str] = set() local_records = dict(context["records"]) for record in object_records: key = str(record.get("key") or "").strip() if not key: raise BulkSeedError(f"{object_api_name} record missing key") object_context = {**context, "records": local_records} materialized_match = materialize_value(record.get("match") or {}, object_context) materialized_fields = materialize_value(record.get("fields") or {}, object_context) if not isinstance(materialized_match, dict) or not materialized_match: raise BulkSeedError(f"{key} must have a non-empty match") if not isinstance(materialized_fields, dict): raise BulkSeedError(f"{key} fields must be an object") signature = match_signature(object_api_name, materialized_match) row = rows_by_signature.get(signature) if row is None: row = { "key": key, "object": object_api_name, "id": stable_uuid(object_api_name, key), "match": materialized_match, "fields": materialized_fields, "alias_keys": [key], } rows_by_signature[signature] = row ordered_signatures.append(signature) else: if row["fields"] != materialized_fields: conflicting_signatures.add(signature) row["fields"] = materialized_fields row["alias_keys"].append(key) row_state = state_for_row(row) for alias_key in row["alias_keys"]: local_records[alias_key] = row_state planned_rows = [rows_by_signature[signature] for signature in ordered_signatures] planned_state: dict[str, dict[str, Any]] = {} for row in planned_rows: row_state = state_for_row(row) for alias_key in row["alias_keys"]: planned_state[alias_key] = row_state collapsed_count = sum(max(0, len(row["alias_keys"]) - 1) for row in planned_rows) duplicate_match_group_count = sum( len(row["alias_keys"]) > 1 for row in planned_rows ) return ( planned_rows, planned_state, collapsed_count, duplicate_match_group_count, len(conflicting_signatures), ) def plan_bulk_seed( seed: dict[str, Any], selected_objects: set[str], *, run_id: str = "baseline", dry_run: bool = False, db: DbConfig | None = None, ) -> dict[str, Any]: key_object_map = row_by_key(seed.get("baseline_records") or []) grouped = records_by_object(seed, selected_objects) inputs = build_inputs(seed, run_id) context = { "seed": seed, "run_id": run_id, "inputs": inputs, "lookups": {}, "records": {}, } if dry_run: for lookup_def in seed.get("lookup_records") or []: key = lookup_def.get("key") context["lookups"][key] = { "id": f"dry-run:{key}", "object": lookup_def.get("object"), } else: if db is None: raise BulkSeedError("db config is required outside dry-run") context["lookups"] = resolve_lookup_records(seed, db, context) pending = dict(grouped) planned_by_object: dict[str, list[dict[str, Any]]] = {} layers: list[list[str]] = [] unsupported: list[dict[str, Any]] = [] fallback_objects = unsupported_unselected_objects(seed, selected_objects) collapsed_count = 0 duplicate_match_group_count = 0 conflicting_duplicate_match_group_count = 0 while pending: layer_successes: list[ tuple[str, list[dict[str, Any]], dict[str, dict[str, Any]], int, int, int] ] = [] layer_unresolved: dict[str, str] = {} layer_start_records = dict(context["records"]) for object_api_name, object_records in list(pending.items()): try: ( rows, state, object_collapsed, object_duplicate_groups, object_conflicting_groups, ) = plan_one_object( object_api_name, object_records, {**context, "records": layer_start_records}, ) except MissingReference as exc: dependency_object = key_object_map.get(exc.key) if dependency_object in pending: layer_unresolved[object_api_name] = f"waiting for {dependency_object}.{exc.key}" continue unsupported.append({"object": object_api_name, "reason": str(exc)}) pending.pop(object_api_name) continue except BulkSeedError as exc: unsupported.append({"object": object_api_name, "reason": str(exc)}) pending.pop(object_api_name) continue layer_successes.append( ( object_api_name, rows, state, object_collapsed, object_duplicate_groups, object_conflicting_groups, ) ) if not layer_successes: for object_api_name, reason in layer_unresolved.items(): unsupported.append({"object": object_api_name, "reason": reason}) pending.pop(object_api_name, None) continue if not dry_run: if db is None: raise BulkSeedError("db config is required outside dry-run") with connect(db) as conn: for ( object_api_name, rows, _state, _object_collapsed, _object_duplicate_groups, _object_conflicting_groups, ) in layer_successes: pre_resolve_existing_ids(conn, object_api_name, rows) layer_objects: list[str] = [] for ( object_api_name, rows, state, object_collapsed, object_duplicate_groups, object_conflicting_groups, ) in layer_successes: state = {} for row in rows: row_state = state_for_row(row) for alias_key in row["alias_keys"]: state[alias_key] = row_state planned_by_object[object_api_name] = rows context["records"].update(state) collapsed_count += object_collapsed duplicate_match_group_count += object_duplicate_groups conflicting_duplicate_match_group_count += object_conflicting_groups pending.pop(object_api_name) layer_objects.append(object_api_name) layers.append(layer_objects) return { "inputs": inputs, "lookups": context["lookups"], "records": context["records"], "planned_by_object": planned_by_object, "layers": layers, "unsupported": unsupported, "fallback_objects": fallback_objects, "collapsed_duplicate_records": collapsed_count, "duplicate_match_group_count": duplicate_match_group_count, "conflicting_duplicate_match_group_count": conflicting_duplicate_match_group_count, } def resolve_lookup_records(seed: dict[str, Any], db: DbConfig, base_context: dict[str, Any]) -> dict[str, dict[str, Any]]: lookups: dict[str, dict[str, Any]] = {} context = {**base_context, "lookups": lookups} with connect(db) as conn: for lookup_def in seed.get("lookup_records") or []: key = str(lookup_def.get("key") or "").strip() object_api_name = str(lookup_def.get("object") or "").strip() match_fields = materialize_value(lookup_def.get("match") or {}, context) rows = select_rows_by_match(conn, object_api_name, match_fields) if len(rows) != 1: raise BulkSeedError(f"lookup {key} expected exactly one {object_api_name} row, found {len(rows)}") lookups[key] = {"id": rows[0]["Id"], "object": object_api_name, **rows[0]} context["lookups"] = lookups return lookups def quote_identifier(identifier: str) -> str: return f"`{identifier.replace('`', '``')}`" def resolve_table_name(conn, object_api_name: str) -> str: with conn.cursor() as cur: cur.execute( """ SELECT table_name FROM information_schema.tables WHERE table_schema = DATABASE() AND LOWER(table_name) = LOWER(%s) ORDER BY CASE WHEN table_name = %s THEN 0 ELSE 1 END, table_name LIMIT 1 """, (object_api_name, object_api_name), ) row = cur.fetchone() if not row: raise BulkSeedError(f"missing target table for {object_api_name}") return str(row[0]) def fetch_table_columns(conn, table: str) -> dict[str, str]: with conn.cursor() as cur: cur.execute( """ SELECT column_name FROM information_schema.columns WHERE table_schema = DATABASE() AND table_name = %s ORDER BY ordinal_position """, (table,), ) rows = cur.fetchall() return {str(row[0]).lower(): str(row[0]) for row in rows} def select_expr_for_requested_column(columns: dict[str, str], requested: str) -> str: actual = columns[requested.lower()] expr = quote_identifier(actual) if actual != requested: expr = f"{expr} AS {quote_identifier(requested)}" return expr def select_rows_by_match(conn, object_api_name: str, match_fields: dict[str, Any]) -> list[dict[str, Any]]: pymysql = require_pymysql() table = resolve_table_name(conn, object_api_name) columns = fetch_table_columns(conn, table) select_columns = ["Id", *[name for name in match_fields if name.lower() in columns]] where_parts = [] params: list[Any] = [] for field, value in match_fields.items(): actual = columns.get(field.lower()) if actual is None: raise BulkSeedError(f"{object_api_name}.{field} is missing from table {table}") where_parts.append(f"{quote_identifier(actual)} <=> %s") params.append(value) sql = ( f"SELECT {', '.join(select_expr_for_requested_column(columns, col) for col in select_columns)} " f"FROM {quote_identifier(table)} WHERE {' AND '.join(where_parts)} LIMIT 2" ) with conn.cursor(pymysql.cursors.DictCursor) as cur: cur.execute(sql, params) return list(cur.fetchall()) def normalize_match_key_value(value: Any) -> Any: if isinstance(value, bool): return int(value) if isinstance(value, Decimal): if value == value.to_integral_value(): return int(value) return float(value) return value def match_key(match_fields: dict[str, Any], fields: list[str]) -> str: return json.dumps( [normalize_match_key_value(match_fields.get(field)) for field in fields], sort_keys=False, separators=(",", ":"), ) def pre_resolve_existing_ids(conn, object_api_name: str, rows: list[dict[str, Any]]) -> None: pymysql = require_pymysql() if not rows: return table = resolve_table_name(conn, object_api_name) columns = fetch_table_columns(conn, table) groups: dict[tuple[str, ...], list[dict[str, Any]]] = {} for row in rows: field_tuple = tuple(row["match"].keys()) groups.setdefault(field_tuple, []).append(row) for match_fields, group_rows in groups.items(): for field in match_fields: if field.lower() not in columns: raise BulkSeedError(f"{object_api_name}.{field} is missing from table {table}") existing: dict[str, str] = {} for start in range(0, len(group_rows), 100): chunk = group_rows[start : start + 100] or_parts = [] params: list[Any] = [] for row in chunk: and_parts = [] for field in match_fields: and_parts.append(f"{quote_identifier(columns[field.lower()])} <=> %s") params.append(row["match"].get(field)) or_parts.append(f"({' AND '.join(and_parts)})") select_cols = ["Id", *match_fields] sql = ( f"SELECT {', '.join(select_expr_for_requested_column(columns, col) for col in select_cols)} " f"FROM {quote_identifier(table)} WHERE {' OR '.join(or_parts)}" ) with conn.cursor(pymysql.cursors.DictCursor) as cur: cur.execute(sql, params) for existing_row in cur.fetchall(): key = match_key(existing_row, list(match_fields)) if key in existing: raise BulkSeedError(f"duplicate existing {object_api_name} match for {key}") existing[key] = existing_row["Id"] for row in group_rows: resolved_id = existing.get(match_key(row["match"], list(match_fields))) if resolved_id: row["id"] = resolved_id def resolve_hashing_backend_dir() -> Path | None: candidates = [ os.environ.get("NEXUSCRM_BACKEND_DIR", ""), str(Path(os.environ["NEXUSCRM_WORKTREE"]) / "backend") if os.environ.get("NEXUSCRM_WORKTREE") else "", str(Path(os.environ["NEXUSCRM_SOURCE_REPO"]) / "backend") if os.environ.get("NEXUSCRM_SOURCE_REPO") else "", ] for candidate in candidates: if not candidate: continue path = Path(candidate) if (path / "go.mod").is_file(): return path return None def hash_user_password(password: Any) -> str: raw_password = str(password) cached = PASSWORD_HASH_CACHE.get(raw_password) if cached: return cached backend_dir = resolve_hashing_backend_dir() if backend_dir is None: raise BulkSeedError( "User.Password bulk hashing needs NEXUSCRM_BACKEND_DIR, " "NEXUSCRM_WORKTREE, or NEXUSCRM_SOURCE_REPO" ) with tempfile.TemporaryDirectory(prefix="exp-bulk-bcrypt-") as temp_dir: program_path = Path(temp_dir) / "main.go" program_path.write_text( """package main import ( "fmt" "os" authpkg "github.com/nexuscrm/backend/pkg/auth" ) func main() { if len(os.Args) != 2 { panic("expected password arg") } hashed, err := authpkg.HashPassword(os.Args[1]) if err != nil { panic(err) } fmt.Print(hashed) } """, encoding="utf-8", ) result = subprocess.run( ["go", "run", str(program_path), raw_password], cwd=backend_dir, check=True, capture_output=True, text=True, ) hashed = result.stdout.strip() if not hashed: raise BulkSeedError("NexusCRM password hasher returned an empty hash") PASSWORD_HASH_CACHE[raw_password] = hashed return hashed def normalize_row_for_table( object_api_name: str, row: dict[str, Any], columns: dict[str, str], admin_user_id: str | None, admin_password_hash: str | None = None, ) -> dict[str, Any]: raw = {"Id": row["id"], **row["match"], **row["fields"]} if object_api_name == "User": if raw.get("Password"): if env_bool("EXP_BULK_HASH_USER_PASSWORDS", True): raw["Password"] = hash_user_password(raw["Password"]) elif admin_password_hash: raw["Password"] = admin_password_hash else: raise BulkSeedError("EXP_BULK_HASH_USER_PASSWORDS=0 requires an admin password hash") elif admin_password_hash: raw["Password"] = admin_password_hash now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S") for column in AUDIT_DATETIME_COLUMNS: if column.lower() in columns and column not in raw: raw[column] = now if "IsDeleted".lower() in columns and "IsDeleted" not in raw: raw["IsDeleted"] = False if admin_user_id: for column in AUDIT_USER_COLUMNS: if column.lower() in columns and column not in raw: raw[column] = admin_user_id normalized: dict[str, Any] = {} for field, value in raw.items(): actual = columns.get(field.lower()) if actual is None: raise BulkSeedError(f"{object_api_name}.{field} is missing from target table") normalized[actual] = value return normalized def seed_update_columns_for_row(object_api_name: str, row: dict[str, Any], columns: dict[str, str]) -> set[str]: seed_values = {**row["match"], **row["fields"]} update_columns: set[str] = set() for field in seed_values: actual = columns.get(field.lower()) if actual is None: raise BulkSeedError(f"{object_api_name}.{field} is missing from target table") if actual not in NEVER_UPDATE_COLUMNS: update_columns.add(actual) return update_columns def ordered_insert_columns(row_values: dict[str, Any], table_columns: dict[str, str]) -> tuple[str, ...]: ordered = [actual for actual in table_columns.values() if actual in row_values] extras = [column for column in row_values if column not in ordered] ordered.extend(extras) id_column = table_columns.get("id", "Id") if id_column in ordered: ordered.remove(id_column) ordered.insert(0, id_column) return tuple(ordered) def build_sparse_upsert_batches( normalized_rows: list[dict[str, Any]], update_columns_by_row: list[set[str]], table_columns: dict[str, str], ) -> list[dict[str, Any]]: if len(normalized_rows) != len(update_columns_by_row): raise BulkSeedError("normalized rows and update columns length mismatch") grouped: dict[tuple[tuple[str, ...], tuple[str, ...]], list[dict[str, Any]]] = {} for row_values, seed_update_columns in zip(normalized_rows, update_columns_by_row): insert_columns = ordered_insert_columns(row_values, table_columns) update_columns = tuple( column for column in insert_columns if column in seed_update_columns and column not in NEVER_UPDATE_COLUMNS ) grouped.setdefault((insert_columns, update_columns), []).append(row_values) return [ { "insert_columns": insert_columns, "update_columns": update_columns, "rows": rows, } for (insert_columns, update_columns), rows in grouped.items() ] def apply_object_rows(db: DbConfig, object_api_name: str, rows: list[dict[str, Any]], admin_user_id: str | None) -> dict[str, Any]: started = time.perf_counter() with connect(db) as conn: table = resolve_table_name(conn, object_api_name) pre_resolve_existing_ids(conn, object_api_name, rows) columns = fetch_table_columns(conn, table) admin_password_hash = None if object_api_name == "User": with conn.cursor() as cur: if admin_user_id: cur.execute("SELECT Password FROM `User` WHERE Id = %s LIMIT 1", (admin_user_id,)) else: cur.execute("SELECT Password FROM `User` WHERE Username = %s LIMIT 1", ("admin@example.invalid",)) admin_row = cur.fetchone() if not admin_row or not admin_row[0]: raise BulkSeedError("admin user password hash is missing") admin_password_hash = str(admin_row[0]) normalized_rows = [ normalize_row_for_table(object_api_name, row, columns, admin_user_id, admin_password_hash) for row in rows ] update_columns_by_row = [ seed_update_columns_for_row(object_api_name, row, columns) for row in rows ] upsert_batches = build_sparse_upsert_batches(normalized_rows, update_columns_by_row, columns) with conn.cursor() as cur: for batch in upsert_batches: all_columns = batch["insert_columns"] update_columns = batch["update_columns"] placeholders = ", ".join(["%s"] * len(all_columns)) insert_columns = ", ".join(quote_identifier(column) for column in all_columns) update_clause = ", ".join( f"{quote_identifier(column)} = VALUES({quote_identifier(column)})" for column in update_columns ) or f"{quote_identifier(columns['id'])} = {quote_identifier(columns['id'])}" sql = ( f"INSERT INTO {quote_identifier(table)} ({insert_columns}) " f"VALUES ({placeholders}) " f"ON DUPLICATE KEY UPDATE {update_clause}" ) values = [[row[column] for column in all_columns] for row in batch["rows"]] cur.executemany(sql, values) verify_ids = [row["id"] for row in rows] found_count = 0 for start in range(0, len(verify_ids), 200): chunk = verify_ids[start : start + 200] id_placeholders = ", ".join(["%s"] * len(chunk)) cur.execute( f"SELECT COUNT(*) FROM {quote_identifier(table)} WHERE {quote_identifier(columns['id'])} IN ({id_placeholders})", chunk, ) found_count += int(cur.fetchone()[0]) conn.commit() duration_ms = round((time.perf_counter() - started) * 1000) if found_count != len(rows): raise BulkSeedError(f"{object_api_name} verification found {found_count}/{len(rows)} rows") return { "object": object_api_name, "rows": len(rows), "aliases": sum(len(row["alias_keys"]) for row in rows), "batches": len(upsert_batches), "duration_ms": duration_ms, } def refresh_state_ids(records_state: dict[str, dict[str, Any]], planned_by_object: dict[str, list[dict[str, Any]]]) -> None: for rows in planned_by_object.values(): for row in rows: row_state = state_for_row(row) for alias_key in row["alias_keys"]: records_state[alias_key] = row_state def apply_plan(db: DbConfig, plan: dict[str, Any], max_workers: int) -> list[dict[str, Any]]: admin_user_id = plan.get("lookups", {}).get("watchdog_admin_user", {}).get("id") results: list[dict[str, Any]] = [] planned_by_object = plan["planned_by_object"] for layer in plan["layers"]: with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, max_workers)) as executor: futures = { executor.submit(apply_object_rows, db, object_api_name, planned_by_object[object_api_name], admin_user_id): object_api_name for object_api_name in layer } for future in concurrent.futures.as_completed(futures): results.append(future.result()) refresh_state_ids(plan["records"], planned_by_object) return results def write_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") def main() -> int: args = parse_args() started = time.perf_counter() seed = load_seed_payload(args) if args.dry_run: os.environ.setdefault("EXP_SEED_APPLY_PASSWORD", "dry-run-password") selected_objects = set(parse_list(args.objects, DEFAULT_BULK_OBJECTS)) selected_record_keys = set(parse_list(args.record_keys, ())) if selected_record_keys: seed = seed_payload_for_selected_record_keys(seed, selected_record_keys) if args.selected_only: if not args.objects: raise BulkSeedError("--selected-only requires --objects") seed = seed_payload_for_selected_objects(seed, selected_objects) db = None if args.dry_run else db_config_from_args(args) plan = plan_bulk_seed(seed, selected_objects, dry_run=args.dry_run, db=db) coverage_complete = not plan["unsupported"] and not plan["fallback_objects"] if not args.dry_run and not coverage_complete: raise BulkSeedError( "bulk seed coverage is incomplete; refusing to apply a partial baseline" ) apply_results: list[dict[str, Any]] = [] if not args.dry_run: assert db is not None apply_results = apply_plan(db, plan, args.max_workers) bulk_objects = [object_api_name for layer in plan["layers"] for object_api_name in layer] bulk_record_count = sum(len(rows) for rows in plan["planned_by_object"].values()) bulk_alias_count = sum( len(row["alias_keys"]) for rows in plan["planned_by_object"].values() for row in rows ) summary = { "dry_run": args.dry_run, "bulk_objects": bulk_objects, "bulk_record_count": bulk_record_count, "bulk_alias_count": bulk_alias_count, "total_alias_count": len(seed.get("baseline_records") or []), "fallback_alias_count": sum(int(item.get("records", 0)) for item in plan["fallback_objects"]), "collapsed_duplicate_records": plan["collapsed_duplicate_records"], "duplicate_match_group_count": plan["duplicate_match_group_count"], "conflicting_duplicate_match_group_count": plan[ "conflicting_duplicate_match_group_count" ], "lookup_record_count": len(seed.get("lookup_records") or []), "app_shell_reference_count": len(seed.get("app_shells") or []), "layers": plan["layers"], "unsupported": plan["unsupported"], "fallback_objects": plan["fallback_objects"], "coverage_complete": coverage_complete, "apply_results": apply_results, "duration_ms": round((time.perf_counter() - started) * 1000), } write_json(Path(args.state_file), {"records": plan["records"]}) write_json(Path(args.summary_file), summary) print(json.dumps(summary, indent=2, sort_keys=True)) return 0 if __name__ == "__main__": raise SystemExit(main())