#!/usr/bin/env python3 from __future__ import annotations import argparse import json import os import sys from pathlib import Path from typing import Any from common import DeployError, db_exec, db_query_scalar, load_env_file from salesforce_file_projection_audit import ( DEFAULT_ATTACHMENT_KIND, DEFAULT_PARENT_OBJECT_API_NAME, DEFAULT_SOURCE_TYPE, DEFAULT_STATUS, DEFAULT_STORAGE_BUCKET, DEFAULT_STORAGE_KEY_PREFIX, DEFAULT_STORAGE_PROVIDER, DEFAULT_SYSTEM_FILE_TABLE, DEFAULT_TARGET_TABLE, audit_salesforce_file_projection, column_lookup, fetch_table_columns, normalize_logical_storage_key_prefix, normalize_storage_key_prefix, require_test_slot_env, sanitize_identifier, split_csv_list, sql_string_literal, unique_ordered, ) RUNTIME_DIR = Path(__file__).resolve().parent REPO_ROOT = RUNTIME_DIR.parents[2] SALESFORCE_MIGRATION_DIR = REPO_ROOT / "scripts" / "salesforce_migration" if str(SALESFORCE_MIGRATION_DIR) not in sys.path: sys.path.insert(0, str(SALESFORCE_MIGRATION_DIR)) from load_salesforce_snapshot_to_tidb import ( # noqa: E402 basename_sql_expr, build_salesforce_file_projection_count_sql, column_expr, mime_type_case_expr, nullish_column_expr, salesforce_file_projection_from_sql, ) DEFAULT_CONTENT_DOCUMENT_LINK_TABLE = "contentdocumentlink" DEFAULT_CONTENT_VERSION_TABLE = "contentversion" SYSTEM_FILE_BASE_INSERT_COLUMNS = ( "Id", "Name", "OriginalFileName", "ParentId", "ParentObjectApiName", "AttachmentKind", "MimeType", "Extension", "SizeBytes", "SourceType", "SourceContentDocumentId", "SourceContentVersionId", "Status", "StorageProvider", "StorageBucket", "StorageKey", "StoragePath", "CreatedDate", "LastModifiedDate", ) SYSTEM_FILE_AUDIT_USER_COLUMNS = ("CreatedById", "LastModifiedById", "OwnerId") SYSTEM_FILE_AUDIT_USER_UPDATE_COLUMNS = ("LastModifiedById", "OwnerId") def query_int(envv: dict[str, str], sql: str) -> int: raw = db_query_scalar(envv, sql, database=envv["TIDB_DATABASE"]) first = raw.splitlines()[0].strip() if raw.splitlines() else "0" return int(first or "0") def missing_tables(table_columns: dict[str, list[str]], table_names: list[str]) -> list[str]: return [name for name in table_names if not table_columns.get(name.lower())] def required_columns_missing(columns: list[str], required: set[str]) -> list[str]: available = {column.lower() for column in columns} return sorted(column for column in required if column.lower() not in available) def system_file_audit_user_columns(system_file_columns: list[str]) -> list[str]: lookup = column_lookup(system_file_columns) return [ lookup[column.lower()] for column in SYSTEM_FILE_AUDIT_USER_COLUMNS if column.lower() in lookup ] def audit_user_sql_expr(user_columns: list[str]) -> str: lookup = column_lookup(user_columns) id_column = lookup.get("id", "Id") if "username" in lookup: username_column = lookup["username"] return ( "COALESCE(" f"(SELECT `{sanitize_identifier(id_column)}` FROM `User` " f"WHERE `{sanitize_identifier(username_column)}` = 'admin@example.invalid' LIMIT 1), " f"(SELECT `{sanitize_identifier(id_column)}` FROM `User` " f"ORDER BY `{sanitize_identifier(id_column)}` LIMIT 1)" ")" ) return ( f"(SELECT `{sanitize_identifier(id_column)}` FROM `User` " f"ORDER BY `{sanitize_identifier(id_column)}` LIMIT 1)" ) def build_raw_projection_insert_sql( *, parent: dict[str, str], content_document_link_table: str, content_version_table: str, system_file_table: str, cdl_columns: list[str], cv_columns: list[str], system_file_columns: list[str], source_type: str, storage_provider: str, storage_bucket: str, storage_key_prefix: str, logical_storage_key_prefix: str = "", audit_user_sql: str = "", ) -> str: system_lookup = column_lookup(system_file_columns) insert_columns = [ system_lookup.get(column.lower(), column) for column in SYSTEM_FILE_BASE_INSERT_COLUMNS ] path_file_name = basename_sql_expr(nullish_column_expr("cv", cv_columns, "PathOnClient")) file_name = ( "COALESCE(" f"{path_file_name}, " f"{nullish_column_expr('cv', cv_columns, 'Title')}, " "cv.`Id`)" ) size_expr = f"COALESCE({column_expr('cv', cv_columns, 'ContentSize', 'NULL')}, 0)" created_expr = ( "COALESCE(" f"{column_expr('cv', cv_columns, 'CreatedDate', 'NULL')}, " f"{column_expr('cdl', cdl_columns, 'CreatedDate', 'NULL')}, " "UTC_TIMESTAMP())" ) modified_expr = ( "COALESCE(" f"{column_expr('cv', cv_columns, 'LastModifiedDate', 'NULL')}, " f"{column_expr('cdl', cdl_columns, 'LastModifiedDate', 'NULL')}, " "UTC_TIMESTAMP())" ) extension_expr = nullish_column_expr("cv", cv_columns, "FileExtension") key_prefix = normalize_storage_key_prefix(storage_key_prefix) logical_key_prefix = normalize_logical_storage_key_prefix( logical_storage_key_prefix ) storage_key_expr = ( f"CONCAT({sql_string_literal(logical_key_prefix + '/' if logical_key_prefix else '')}, " "cv.`Id`)" ) storage_path_expr = ( f"CONCAT('s3://', {sql_string_literal(storage_bucket)}, '/', " f"{sql_string_literal(key_prefix + '/')}, cv.`Id`)" ) select_exprs = [ "CONCAT('sfcdl-', cdl.`Id`)", file_name, file_name, "parent_record.`Id`", sql_string_literal(parent["parent_object_api_name"]), sql_string_literal(parent["attachment_kind"]), mime_type_case_expr(cv_columns), extension_expr, size_expr, sql_string_literal(source_type), "cdl.`ContentDocumentId`", "cv.`Id`", "'Active'", sql_string_literal(storage_provider), sql_string_literal(storage_bucket), storage_key_expr, storage_path_expr, created_expr, modified_expr, ] if audit_user_sql: inserted_lower = {column.lower() for column in insert_columns} for column in system_file_audit_user_columns(system_file_columns): if column.lower() in inserted_lower: continue insert_columns.append(column) select_exprs.append(audit_user_sql) inserted_lower.add(column.lower()) insert_columns_sql = ",\n".join(f" `{sanitize_identifier(column)}`" for column in insert_columns) select_sql = ",\n".join( f" {expr} AS `{sanitize_identifier(column)}`" for column, expr in zip(insert_columns, select_exprs) ) update_columns = [ "Name", "OriginalFileName", "ParentId", "ParentObjectApiName", "AttachmentKind", "MimeType", "Extension", "SizeBytes", "SourceType", "SourceContentDocumentId", "SourceContentVersionId", "Status", "StorageProvider", "StorageBucket", "StorageKey", "StoragePath", "LastModifiedDate", ] present_insert_columns = {column.lower(): column for column in insert_columns} for column in SYSTEM_FILE_AUDIT_USER_UPDATE_COLUMNS: if column.lower() in present_insert_columns: update_columns.append(present_insert_columns[column.lower()]) update_sql = ",\n".join( f" `{sanitize_identifier(column)}` = VALUES(`{sanitize_identifier(column)}`)" for column in update_columns ) from_sql = salesforce_file_projection_from_sql( parent=parent, content_document_link_table=content_document_link_table, content_version_table=content_version_table, cdl_columns=cdl_columns, cv_columns=cv_columns, ) return f""" INSERT INTO `{sanitize_identifier(system_file_table)}` ( {insert_columns_sql} ) SELECT {select_sql} {from_sql} ON DUPLICATE KEY UPDATE {update_sql} """.strip() def load_required_table_columns( envv: dict[str, str], table_names: list[str], ) -> dict[str, list[str]]: return {name.lower(): fetch_table_columns(envv, name) for name in table_names} def repair_salesforce_file_projection_from_raw( envv: dict[str, str], *, apply: bool, target_parent_keys: list[str], parent_object_api_name: str = DEFAULT_PARENT_OBJECT_API_NAME, target_table: str = DEFAULT_TARGET_TABLE, system_file_table: str = DEFAULT_SYSTEM_FILE_TABLE, content_document_link_table: str = DEFAULT_CONTENT_DOCUMENT_LINK_TABLE, content_version_table: str = DEFAULT_CONTENT_VERSION_TABLE, attachment_kind: str = DEFAULT_ATTACHMENT_KIND, source_type: str = DEFAULT_SOURCE_TYPE, storage_provider: str = DEFAULT_STORAGE_PROVIDER, storage_bucket: str = DEFAULT_STORAGE_BUCKET, storage_key_prefix: str = DEFAULT_STORAGE_KEY_PREFIX, logical_storage_key_prefix: str = "", status: str = DEFAULT_STATUS, allow_non_test_env: bool = False, ) -> dict[str, Any]: require_test_slot_env(envv, allow_non_test_env=allow_non_test_env) target_keys = unique_ordered(target_parent_keys) if not target_keys: raise DeployError("at least one target parent key is required before raw repair") before = audit_salesforce_file_projection( envv, target_parent_keys=target_keys, parent_object_api_name=parent_object_api_name, target_table=target_table, system_file_table=system_file_table, attachment_kind=attachment_kind, source_type=source_type, storage_provider=storage_provider, storage_bucket=storage_bucket, storage_key_prefix=storage_key_prefix, logical_storage_key_prefix=logical_storage_key_prefix, status=status, allow_non_test_env=True, ) summary: dict[str, Any] = { "status": before.get("status", "unknown"), "mode": "apply" if apply else "dry-run", "env_name": envv.get("ENV_NAME", ""), "public_base_url": envv.get("PUBLIC_BASE_URL", ""), "filters": { "parent_object_api_name": parent_object_api_name, "target_table": sanitize_identifier(target_table), "system_file_table": sanitize_identifier(system_file_table), "content_document_link_table": sanitize_identifier(content_document_link_table), "content_version_table": sanitize_identifier(content_version_table), "attachment_kind": attachment_kind, "source_type": source_type, "status": status, "storage_provider": storage_provider, "storage_bucket": storage_bucket, "storage_key_prefix": normalize_storage_key_prefix(storage_key_prefix), "logical_storage_key_prefix": normalize_logical_storage_key_prefix( logical_storage_key_prefix ), }, "before": before, "repair": { "source": "raw-salesforce-tables", "applied": False, "skipped_reason": "", "projectable_rows": 0, }, } resolved_target_keys = { str(target.get("target_key") or "").strip() for target in (before.get("targets") or []) if isinstance(target, dict) and target.get("target_found") and str(target.get("parent_id") or "").strip() } unresolved_target_keys = [ target_key for target_key in target_keys if target_key not in resolved_target_keys ] if unresolved_target_keys: summary["status"] = "blocked_target_resolution" summary["repair"]["skipped_reason"] = "target_parent_keys_not_exact" summary["repair"]["unresolved_target_keys"] = unresolved_target_keys return summary if before.get("status") == "passed": summary["repair"]["skipped_reason"] = "targets_already_covered" return summary table_names = [ target_table, system_file_table, content_document_link_table, content_version_table, "User", ] table_columns = load_required_table_columns(envv, table_names) raw_missing_tables = missing_tables( table_columns, [target_table, system_file_table, content_document_link_table, content_version_table], ) if raw_missing_tables: summary["status"] = "missing_raw_tables" summary["repair"]["skipped_reason"] = "missing_raw_tables" summary["repair"]["missing_tables"] = raw_missing_tables return summary cdl_columns = table_columns[content_document_link_table.lower()] cv_columns = table_columns[content_version_table.lower()] missing_columns = { content_document_link_table: required_columns_missing( cdl_columns, {"Id", "ContentDocumentId", "LinkedEntityId"}, ), content_version_table: required_columns_missing( cv_columns, {"Id", "ContentDocumentId"}, ), target_table: required_columns_missing(table_columns[target_table.lower()], {"Id"}), } missing_columns = {table: columns for table, columns in missing_columns.items() if columns} if missing_columns: summary["status"] = "missing_raw_columns" summary["repair"]["skipped_reason"] = "missing_raw_columns" summary["repair"]["missing_columns"] = missing_columns return summary parent = { "parent_object_api_name": parent_object_api_name, "target_table": sanitize_identifier(target_table), "attachment_kind": attachment_kind, } count_sql = build_salesforce_file_projection_count_sql( parent=parent, content_document_link_table=content_document_link_table, content_version_table=content_version_table, cdl_columns=cdl_columns, cv_columns=cv_columns, require_version=True, ) projectable_rows = query_int(envv, count_sql) summary["repair"]["projectable_rows"] = projectable_rows if projectable_rows <= 0: summary["status"] = "empty_raw_projection" summary["repair"]["skipped_reason"] = "empty_raw_projection" return summary if not apply: summary["status"] = "would_repair" summary["repair"]["skipped_reason"] = "dry_run" return summary audit_user_sql = "" audit_columns = system_file_audit_user_columns(table_columns[system_file_table.lower()]) if audit_columns: audit_user_sql = audit_user_sql_expr(table_columns.get("user", [])) summary["repair"]["system_file_audit_user_columns"] = audit_columns insert_sql = build_raw_projection_insert_sql( parent=parent, content_document_link_table=content_document_link_table, content_version_table=content_version_table, system_file_table=system_file_table, cdl_columns=cdl_columns, cv_columns=cv_columns, system_file_columns=table_columns[system_file_table.lower()], source_type=source_type, storage_provider=storage_provider, storage_bucket=storage_bucket, storage_key_prefix=storage_key_prefix, logical_storage_key_prefix=logical_storage_key_prefix, audit_user_sql=audit_user_sql, ) db_exec(envv, insert_sql, database=envv["TIDB_DATABASE"]) summary["repair"]["applied"] = True after = audit_salesforce_file_projection( envv, target_parent_keys=target_keys, parent_object_api_name=parent_object_api_name, target_table=target_table, system_file_table=system_file_table, attachment_kind=attachment_kind, source_type=source_type, storage_provider=storage_provider, storage_bucket=storage_bucket, storage_key_prefix=storage_key_prefix, logical_storage_key_prefix=logical_storage_key_prefix, status=status, allow_non_test_env=True, ) summary["after"] = after summary["status"] = "passed" if after.get("status") == "passed" else "failed" return summary def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Repair Salesforce migrated file projection from already-loaded raw tables." ) parser.add_argument("--env-file", required=True) parser.add_argument("--parent-object-api-name", default=DEFAULT_PARENT_OBJECT_API_NAME) parser.add_argument("--target-table", default=DEFAULT_TARGET_TABLE) parser.add_argument("--system-file-table", default=DEFAULT_SYSTEM_FILE_TABLE) parser.add_argument("--content-document-link-table", default=DEFAULT_CONTENT_DOCUMENT_LINK_TABLE) parser.add_argument("--content-version-table", default=DEFAULT_CONTENT_VERSION_TABLE) parser.add_argument("--attachment-kind", default=DEFAULT_ATTACHMENT_KIND) parser.add_argument("--source-type", default=DEFAULT_SOURCE_TYPE) parser.add_argument("--storage-provider", default=DEFAULT_STORAGE_PROVIDER) parser.add_argument("--storage-bucket", default=os.environ.get("SF_S3_BUCKET", DEFAULT_STORAGE_BUCKET)) parser.add_argument("--storage-key-prefix", default=DEFAULT_STORAGE_KEY_PREFIX) parser.add_argument("--logical-storage-key-prefix", default="") parser.add_argument("--status", default=DEFAULT_STATUS) parser.add_argument( "--target-parent-key", action="append", default=[], help="Parent record key to prove after repair. May be repeated or comma-separated.", ) parser.add_argument("--apply", action="store_true") parser.add_argument("--allow-non-test-env", action="store_true") parser.add_argument("--summary-json") parser.add_argument("--fail-on-unrepaired", action="store_true") return parser.parse_args() def main() -> int: args = parse_args() envv = load_env_file(Path(args.env_file).expanduser().resolve()) summary = repair_salesforce_file_projection_from_raw( envv, apply=args.apply, target_parent_keys=[ item for raw in args.target_parent_key for item in split_csv_list(raw) ], parent_object_api_name=args.parent_object_api_name, target_table=args.target_table, system_file_table=args.system_file_table, content_document_link_table=args.content_document_link_table, content_version_table=args.content_version_table, attachment_kind=args.attachment_kind, source_type=args.source_type, storage_provider=args.storage_provider, storage_bucket=args.storage_bucket, storage_key_prefix=args.storage_key_prefix, logical_storage_key_prefix=args.logical_storage_key_prefix, status=args.status, allow_non_test_env=args.allow_non_test_env, ) output = json.dumps(summary, indent=2, ensure_ascii=False) if args.summary_json: Path(args.summary_json).write_text(output + "\n", encoding="utf-8") print(output) if args.fail_on_unrepaired and summary.get("status") != "passed": return 1 return 0 if __name__ == "__main__": raise SystemExit(main())