#!/usr/bin/env python3 from __future__ import annotations import argparse import json import os import sys from pathlib import Path from typing import Any, Callable from common import DeployError, 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, normalize_logical_storage_key_prefix, normalize_storage_key_prefix, require_test_slot_env, sanitize_identifier, split_csv_list, 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 DbConfig, build_export_object_map, collect_manifest_objects, derive_content_blob_upload_state, extract_s3_bucket, load_export_config, manifest_file_path, run_salesforce_file_projection_from_snapshot, ) DEFAULT_CONTENT_DOCUMENT_LINK_TABLE = "contentdocumentlink" DEFAULT_CONTENT_VERSION_TABLE = "contentversion" DEFAULT_SNAPSHOT_MANIFEST_NAME = "manifest.json" DEFAULT_EXPORT_CONFIG = SALESFORCE_MIGRATION_DIR / "configs" / "daily_full_export.json" DEFAULT_SNAPSHOT_SEARCH_ROOTS = ( Path("/data/sf-export"), Path("/home/ec2-user/sf-export"), ) FALLBACK_SNAPSHOT_MANIFEST_NAMES = ("manifest.json", "manifest_phase1_source_snapshot.json") def db_config_from_env(envv: dict[str, str]) -> DbConfig: missing = [ key for key in ["TIDB_HOST", "TIDB_PORT", "TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"] if not str(envv.get(key) or "").strip() ] if missing: raise DeployError(f"missing TiDB env keys: {', '.join(missing)}") return DbConfig( host=str(envv["TIDB_HOST"]).strip(), port=int(str(envv["TIDB_PORT"]).strip()), user=str(envv["TIDB_USER"]).strip(), password=str(envv["TIDB_PASSWORD"]).strip(), database=str(envv["TIDB_DATABASE"]).strip(), ssl_mode=str(envv.get("TIDB_SSL_MODE") or "REQUIRED").strip() or "REQUIRED", ) def split_path_list(raw: str) -> list[str]: normalized = str(raw or "").replace(",", os.pathsep) return [item.strip() for item in normalized.split(os.pathsep) if item.strip()] def snapshot_manifest_names(preferred: str) -> list[str]: names: list[str] = [] for name in [preferred, *FALLBACK_SNAPSHOT_MANIFEST_NAMES]: normalized = str(name or "").strip() if normalized and normalized not in names: names.append(normalized) return names def first_explicit_snapshot_dir( *, snapshot_dir: str, envv: dict[str, str], ) -> tuple[str, str]: candidates = [ ("argument", snapshot_dir), ("env_file", str(envv.get("SF_SNAPSHOT_DIR") or "")), ("environment", os.environ.get("SF_SNAPSHOT_DIR", "")), ("environment", os.environ.get("SALESFORCE_SNAPSHOT_DIR", "")), ] for source, raw in candidates: text = str(raw or "").strip() if text: return source, text return "", "" def automatic_snapshot_candidates( *, envv: dict[str, str], snapshot_search_roots: list[str], ) -> list[Path]: roots: list[Path] = [] for raw in snapshot_search_roots: roots.extend(Path(item).expanduser() for item in split_path_list(raw)) for raw in ( str(envv.get("SF_SNAPSHOT_SEARCH_ROOTS") or ""), os.environ.get("SF_SNAPSHOT_SEARCH_ROOTS", ""), ): roots.extend(Path(item).expanduser() for item in split_path_list(raw)) if not roots: roots.extend(DEFAULT_SNAPSHOT_SEARCH_ROOTS) candidates: list[Path] = [] for root in roots: if not root.exists(): continue candidates.append(root) candidates.extend(path for path in root.glob("*/*") if path.is_dir()) seen: set[Path] = set() unique: list[Path] = [] for candidate in candidates: resolved = candidate.resolve() if resolved in seen: continue seen.add(resolved) unique.append(resolved) unique.sort(key=lambda path: path.stat().st_mtime if path.exists() else 0, reverse=True) return unique def collect_snapshot_objects( *, snapshot_dir: Path, manifest_name: str, export_config_path: Path, ) -> tuple[list[dict[str, Any]], dict[str, Any] | None]: try: export_config = load_export_config(export_config_path) export_object_map = build_export_object_map(export_config) return collect_manifest_objects(snapshot_dir, manifest_name, export_object_map), None except Exception as exc: return [], {"status": "invalid_manifest", "error": str(exc)} def snapshot_has_required_files( *, snapshot_dir: Path, manifest_objects: list[dict[str, Any]], ) -> list[str]: missing = [] for object_name in ("ContentDocumentLink", "ContentVersion"): path = manifest_file_path(snapshot_dir, manifest_objects, object_name) if path is None or not path.is_file(): missing.append(object_name) return missing def resolve_snapshot_source( *, snapshot_dir: str, snapshot_manifest_name: str, export_config_path: Path, snapshot_search_roots: list[str], envv: dict[str, str], ) -> tuple[Path | None, str, list[dict[str, Any]], dict[str, Any]]: sources, meta = resolve_snapshot_source_candidates( snapshot_dir=snapshot_dir, snapshot_manifest_name=snapshot_manifest_name, export_config_path=export_config_path, snapshot_search_roots=snapshot_search_roots, envv=envv, ) if not sources: return None, "", [], meta return sources[0] def resolve_snapshot_source_candidates( *, snapshot_dir: str, snapshot_manifest_name: str, export_config_path: Path, snapshot_search_roots: list[str], envv: dict[str, str], ) -> tuple[list[tuple[Path, str, list[dict[str, Any]], dict[str, Any]]], dict[str, Any]]: explicit_source, explicit_dir = first_explicit_snapshot_dir( snapshot_dir=snapshot_dir, envv=envv, ) meta: dict[str, Any] = { "requested_dir": explicit_dir, "requested_source": explicit_source, "search_roots": snapshot_search_roots or [str(path) for path in DEFAULT_SNAPSHOT_SEARCH_ROOTS], "checked_count": 0, } if explicit_dir: candidates = [Path(explicit_dir).expanduser()] meta["explicit"] = True else: candidates = automatic_snapshot_candidates( envv=envv, snapshot_search_roots=snapshot_search_roots, ) meta["explicit"] = False if not candidates: meta["status"] = "snapshot_not_found" return [], meta checked_samples: list[str] = [] sources: list[tuple[Path, str, list[dict[str, Any]], dict[str, Any]]] = [] for candidate in candidates: candidate = candidate.resolve() if len(checked_samples) < 8: checked_samples.append(str(candidate)) if not candidate.exists() or not candidate.is_dir(): continue for manifest_name in snapshot_manifest_names(snapshot_manifest_name): manifest_path = candidate / manifest_name if not manifest_path.is_file(): continue meta["checked_count"] = int(meta["checked_count"]) + 1 manifest_objects, error_meta = collect_snapshot_objects( snapshot_dir=candidate, manifest_name=manifest_name, export_config_path=export_config_path, ) if error_meta: meta.update(error_meta) meta["candidate_dir"] = str(candidate) meta["manifest_name"] = manifest_name continue missing_files = snapshot_has_required_files( snapshot_dir=candidate, manifest_objects=manifest_objects, ) if missing_files: meta["status"] = "missing_snapshot_files" meta["candidate_dir"] = str(candidate) meta["manifest_name"] = manifest_name meta["missing_files"] = missing_files continue candidate_meta = { **meta, "status": "found", "candidate_dir": str(candidate), "manifest_name": manifest_name, "manifest_object_count": len(manifest_objects), } candidate_meta.pop("missing_files", None) sources.append((candidate, manifest_name, manifest_objects, candidate_meta)) if explicit_dir: return sources, candidate_meta if sources: meta = {**sources[0][3], "candidate_count": len(sources)} return sources, meta meta.setdefault("status", "snapshot_not_found") meta["checked_samples"] = checked_samples return [], meta def snapshot_storage( *, manifest_objects: list[dict[str, Any]], storage_provider: str, fallback_storage_bucket: str, fallback_storage_key_prefix: str, prefer_manifest_storage: bool = True, ) -> dict[str, str]: content_version_manifest = next( (obj for obj in manifest_objects if str(obj.get("object") or "") == "ContentVersion"), {}, ) bucket = ( extract_s3_bucket(content_version_manifest.get("metadata_s3_uri")) if prefer_manifest_storage else "" ) or fallback_storage_bucket key_prefix = normalize_storage_key_prefix( ( content_version_manifest.get("blob_prefix") if prefer_manifest_storage else "" ) or fallback_storage_key_prefix ) return { "provider": storage_provider, "bucket": bucket, "key_prefix": key_prefix, } def repair_salesforce_file_projection_from_snapshot( 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, snapshot_dir: str = "", snapshot_manifest_name: str = DEFAULT_SNAPSHOT_MANIFEST_NAME, export_config: str = str(DEFAULT_EXPORT_CONFIG), snapshot_search_roots: list[str] | None = None, allow_non_test_env: bool = False, prefer_snapshot_storage: bool = True, create_only: bool = False, expected_parent_ids: list[str] | None = None, skip_if_targets_covered: bool = True, before_projection_commit: Callable[[], None] | None = None, ) -> dict[str, Any]: require_test_slot_env(envv, allow_non_test_env=allow_non_test_env) target_keys = unique_ordered(target_parent_keys) 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), }, "before": before, "repair": { "source": "local-salesforce-snapshot", "applied": False, "skipped_reason": "", "projectable_rows": 0, }, } if before.get("status") == "passed" and skip_if_targets_covered: summary["repair"]["skipped_reason"] = "targets_already_covered" return summary audit_targets = before.get("targets") or [] resolved_parent_ids = unique_ordered( [ str(target.get("parent_id") or "").strip() for target in audit_targets if isinstance(target, dict) and target.get("target_found") ] ) if len(resolved_parent_ids) != len(target_keys): summary["status"] = "blocked_target_resolution" summary["repair"]["skipped_reason"] = "target_parent_ids_not_exact" summary["repair"]["resolved_parent_ids"] = resolved_parent_ids return summary summary["repair"]["target_parent_ids"] = resolved_parent_ids bound_parent_ids = unique_ordered(expected_parent_ids or []) if bound_parent_ids and resolved_parent_ids != bound_parent_ids: summary["status"] = "blocked_target_resolution_drift" summary["repair"]["skipped_reason"] = "target_parent_ids_changed" summary["repair"]["expected_parent_ids"] = bound_parent_ids return summary snapshot_sources, snapshot_meta = resolve_snapshot_source_candidates( snapshot_dir=snapshot_dir, snapshot_manifest_name=snapshot_manifest_name, export_config_path=Path(export_config).expanduser().resolve(), snapshot_search_roots=snapshot_search_roots or [], envv=envv, ) summary["snapshot"] = snapshot_meta if not snapshot_sources: summary["status"] = snapshot_meta.get("status", "snapshot_not_found") summary["repair"]["skipped_reason"] = summary["status"] return summary parent = { "parent_object_api_name": parent_object_api_name, "target_table": sanitize_identifier(target_table), "attachment_kind": attachment_kind, } attempts: list[dict[str, Any]] = [] summary["snapshot_attempts"] = attempts for resolved_dir, resolved_manifest_name, manifest_objects, candidate_meta in snapshot_sources: blob_state = derive_content_blob_upload_state(manifest_objects) candidate_meta = {**candidate_meta, "content_blob_upload_state": blob_state} attempt: dict[str, Any] = { "candidate_dir": str(resolved_dir), "manifest_name": resolved_manifest_name, "status": "pending", } attempts.append(attempt) summary["snapshot"] = candidate_meta if not blob_state.get("complete", True): attempt["status"] = "blocked_incomplete_content_blobs" if candidate_meta.get("explicit"): summary["status"] = "blocked_incomplete_content_blobs" summary["repair"]["skipped_reason"] = "incomplete_content_blobs" return summary continue snapshot_source_storage = snapshot_storage( manifest_objects=manifest_objects, storage_provider=storage_provider, fallback_storage_bucket=storage_bucket, fallback_storage_key_prefix=storage_key_prefix, ) storage = snapshot_storage( manifest_objects=manifest_objects, storage_provider=storage_provider, fallback_storage_bucket=storage_bucket, fallback_storage_key_prefix=storage_key_prefix, prefer_manifest_storage=prefer_snapshot_storage, ) storage["logical_key_prefix"] = normalize_logical_storage_key_prefix( logical_storage_key_prefix ) summary["repair"]["snapshot_source_storage"] = snapshot_source_storage attempt["storage_bucket"] = storage["bucket"] attempt["storage_key_prefix"] = storage["key_prefix"] summary["repair"]["storage_provider"] = storage["provider"] summary["repair"]["storage_bucket"] = storage["bucket"] summary["repair"]["storage_key_prefix"] = storage["key_prefix"] summary["repair"]["logical_storage_key_prefix"] = storage[ "logical_key_prefix" ] if not storage["bucket"]: attempt["status"] = "missing_snapshot_storage_bucket" if candidate_meta.get("explicit"): summary["status"] = "missing_snapshot_storage_bucket" summary["repair"]["skipped_reason"] = "missing_snapshot_storage_bucket" return summary continue try: result = run_salesforce_file_projection_from_snapshot( db=db_config_from_env(envv), snapshot_dir=resolved_dir, manifest_objects=manifest_objects, config={ "parents": [parent], "content_document_link_table": content_document_link_table, "content_version_table": content_version_table, "system_file_table": system_file_table, "source_type": source_type, }, storage=storage, fallback_reason=str(before.get("status") or "audit_failed"), replace_existing=False, blob_state=blob_state, target_parent_ids=resolved_parent_ids, dry_run=not apply, create_only=create_only, before_commit=before_projection_commit, ) except RuntimeError as exc: attempt["status"] = "blocked_target_projection" attempt["error"] = str(exc)[:1000] if candidate_meta.get("explicit"): summary["status"] = "blocked_target_projection" summary["repair"]["skipped_reason"] = ( "target_snapshot_projection_incomplete" ) return summary continue summary["repair"]["applied"] = apply summary["repair"]["projectable_rows"] = int(result.get("projectable_rows") or 0) summary["repair"]["affected_rows"] = int(result.get("affected_rows") or 0) summary["repair"]["snapshot_result"] = result attempt["status"] = str(result.get("status") or "projected_from_snapshot") attempt["projectable_rows"] = int(result.get("projectable_rows") or 0) attempt["affected_rows"] = int(result.get("affected_rows") or 0) parent_results = result.get("parents") or [] if parent_results: attempt["parents"] = parent_results if not apply: if int(result.get("projectable_rows") or 0) <= 0: summary["status"] = "blocked_empty_target_projection" summary["repair"]["skipped_reason"] = "empty_target_projection" return summary summary["status"] = "would_repair" summary["repair"]["skipped_reason"] = "dry_run" attempt["status"] = "would_repair" return summary 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=storage["logical_key_prefix"], status=status, allow_non_test_env=True, ) summary["after"] = after attempt["after_status"] = after.get("status") attempt["after_target_keys_with_files"] = after.get("target_keys_with_files") if after.get("status") == "passed": summary["status"] = "passed" return summary summary["status"] = "failed" if not summary["repair"].get("applied"): summary["repair"]["skipped_reason"] = "no_applicable_snapshot_candidate" return summary def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Repair Salesforce migrated file projection from a local snapshot." ) 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("--status", default=DEFAULT_STATUS) parser.add_argument("--snapshot-dir", default=os.environ.get("SF_SNAPSHOT_DIR", "")) parser.add_argument("--snapshot-manifest-name", default=DEFAULT_SNAPSHOT_MANIFEST_NAME) parser.add_argument("--export-config", default=str(DEFAULT_EXPORT_CONFIG)) parser.add_argument( "--snapshot-search-root", action="append", default=[], help="Root to scan for local Salesforce snapshots when --snapshot-dir is not provided. May be repeated.", ) 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_snapshot( 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, status=args.status, snapshot_dir=args.snapshot_dir, snapshot_manifest_name=args.snapshot_manifest_name, export_config=args.export_config, snapshot_search_roots=args.snapshot_search_root, 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())