ss-tools/backend/src/core/migration_engine.py

# #region Core.MigrationEngine.MigrationEngineModule [C:5] [TYPE Module] [SEMANTICS migration, superset, archive, migration-engine]
# @defgroup Core Module group.
#
# @BRIEF Transforms Superset export ZIP archives while preserving archive integrity and patching mapped identifiers.
# @LAYER     Domain
# @RELATION DEPENDS_ON -> [Core.Logger.LoggerModule]
# @RELATION DEPENDS_ON -> [Core.MappingService.IdMappingService]
# @RELATION DEPENDS_ON -> [Models.Mapping.ResourceType]
# @RELATION DEPENDS_ON -> [EXT:Library:yaml]
# @PRE  Input archives are readable Superset exports and optional mapping collaborators expose remote id lookup APIs.
# @POST  Migration engine contracts preserve ZIP integrity while exposing transformation entrypoints for import pipelines.
# @SIDE_EFFECT  Reads and writes temporary archive contents during transformation workflows and emits structured belief-state logs.
# @DATA_CONTRACT  Input[zip_path, output_path, db_mapping, target_env_id?, fix_cross_filters?] -> Output[Transformed Superset archive]
# @INVARIANT  ZIP structure and non-targeted metadata must remain valid after transformation.
# @INVARIANT  An import archive contains only mapped database resources referenced by its datasets.
# @RATIONALE  Dedicated module for Superset export ZIP transformation (extract → transform → re-package) because archive manipulation has distinct lifecycle requirements that benefit from isolation from API routing and task orchestration layers.
# @REJECTED  Performing ZIP transformations inline within API route handlers was rejected — it would duplicate extraction/packaging logic across endpoints and make archive corruption handling inconsistent.
import json
import os
from pathlib import Path
import re
import tempfile
import zipfile

import yaml

from src.core.mapping_service import IdMappingService
from src.models.mapping import ResourceType

from .logger import belief_scope, logger


# #region Core.MigrationEngine [TYPE Class]
# @defgroup Core Module group.
# @BRIEF Engine for transforming Superset export ZIPs.
# @RELATION DEPENDS_ON -> [Core.MigrationEngine]
class MigrationEngine:
    # #region Core.MigrationEngine.Init [TYPE Function]
    # @PURPOSE: Initializes migration orchestration dependencies for ZIP/YAML metadata transformations.
    # @PRE  mapping_service is None or implements batch remote ID lookup for ResourceType.CHART.
    # @POST  self.mapping_service is assigned and available for optional cross-filter patching flows.
    # @SIDE_EFFECT  Mutates in-memory engine state by storing dependency reference.
    # @DATA_CONTRACT  Input[Optional[IdMappingService]] -> Output[MigrationEngine]
    # @PARAM  mapping_service (Optional[IdMappingService]) - Used for resolving target environment integer IDs.
    def __init__(self, mapping_service: IdMappingService | None = None):
        with belief_scope("MigrationEngine.__init__"):
            logger.reason("Initializing MigrationEngine")
            self.mapping_service = mapping_service
            logger.reflect("MigrationEngine initialized")
    # #endregion Core.MigrationEngine.Init
    # #region Core.MigrationEngine.TransformZip [TYPE Function]
# @ingroup Core
    # @PURPOSE: Extracts ZIP, replaces database UUIDs in YAMLs, patches cross-filters, and re-packages.
    # @RELATION DEPENDS_ON -> [Core.MigrationEngine]
    # @PARAM    zip_path (str) - Path to the source ZIP file.
    # @PARAM    output_path (str) - Path where the transformed ZIP will be saved.
    # @PARAM    db_mapping (Dict[str, str]) - Mapping of source UUID to target UUID.
    # @PARAM    strip_databases (bool) - Whether to remove the databases directory from the archive.
    # @PARAM    target_env_id (Optional[str]) - Used if fix_cross_filters is True to know which environment map to use.
    # @PARAM    fix_cross_filters (bool) - Whether to patch dashboard json_metadata.
    # @PRE  zip_path points to a readable ZIP; output_path parent is writable; db_mapping keys/values are UUID strings.
    # @POST  Returns True only when extraction, transformation, and packaging complete without exception.
    # @SIDE_EFFECT  Reads/writes filesystem archives, creates temporary directory, emits structured logs.
    # @DATA_CONTRACT  Input[(str zip_path, str output_path, Dict[str,str] db_mapping, bool strip_databases, Optional[str] target_env_id, bool fix_cross_filters)] -> Output[bool]
    # @RETURN  bool - True if successful.
    def transform_zip(
        self,
        zip_path: str,
        output_path: str,
        db_mapping: dict[str, str],
        strip_databases: bool = True,
        target_env_id: str | None = None,
        fix_cross_filters: bool = False,
    ) -> bool:
        """
        Transform a Superset export ZIP by replacing database UUIDs and optionally fixing cross-filters.
        """
        with belief_scope("MigrationEngine.transform_zip"):
            logger.reason(f"Starting ZIP transformation: {zip_path} -> {output_path}")
            with tempfile.TemporaryDirectory() as temp_dir_str:
                temp_dir = Path(temp_dir_str)
                try:
                    # 1. Extract
                    logger.reason(f"Extracting source archive to {temp_dir}")
                    with zipfile.ZipFile(zip_path, "r") as zf:
                        zf.extractall(temp_dir)
                    # 2. Transform YAMLs (Datasets)
                    dataset_files = list(temp_dir.glob("**/datasets/**/*.yaml")) + list(
                        temp_dir.glob("**/datasets/*.yaml")
                    )
                    dataset_files = list(set(dataset_files))
                    logger.reason(
                        f"Transforming {len(dataset_files)} dataset YAML files"
                    )
                    required_source_database_uuids = self._collect_dataset_database_uuids(
                        dataset_files
                    )
                    # @RATIONALE Empty db_mapping means identity UUIDs (replace_db_config=False):
                    #   keep source database UUIDs and let target match by UUID.
                    #   Partial mapping is unsafe (some datasets remapped, others not) → refuse.
                    # @REJECTED Always requiring a full mapping was rejected — it made dry-run and
                    #   execute fail with empty mapping even when replace_db_config=False (ADR-0019).
                    unmapped_database_uuids = (
                        required_source_database_uuids - set(db_mapping)
                    )
                    if db_mapping and unmapped_database_uuids:
                        logger.explore(
                            "Archive references databases without a target mapping; refusing partial mapping",
                            payload={
                                "unmapped_database_count": len(unmapped_database_uuids),
                                "mapped_count": len(db_mapping),
                            },
                        )
                        return False
                    if not db_mapping and required_source_database_uuids:
                        logger.reason(
                            "No DB mapping provided; keeping source database UUIDs (identity)",
                            payload={
                                "database_uuid_count": len(required_source_database_uuids),
                            },
                        )
                    for ds_file in dataset_files:
                        self._transform_yaml(ds_file, db_mapping)
                    # 2.1 Transform YAMLs (Databases — replace UUID with target UUID)
                    #     When a database UUID in the archive matches a target UUID that exists
                    #     in the target Superset, Superset's import_database() will find it,
                    #     skip creation, and populate database_ids — avoiding password errors.
                    db_files = list(temp_dir.glob("**/databases/**/*.yaml")) + list(
                        temp_dir.glob("**/databases/*.yaml")
                    )
                    db_files = list(set(db_files))
                    referenced_database_files: set[Path] = set()
                    if db_files:
                        logger.reason(
                            f"Transforming {len(db_files)} database YAML files"
                        )
                        for db_file in db_files:
                            if self._database_yaml_uuid(db_file) in required_source_database_uuids:
                                self._transform_database_yaml(db_file, db_mapping)
                                referenced_database_files.add(db_file)
                            else:
                                logger.reason(
                                    f"Excluding unreferenced database resource {db_file.name}"
                                )
                    # 2.5 Patch Cross-Filters (Dashboards)
                    if fix_cross_filters:
                        if self.mapping_service and target_env_id:
                            dash_files = list(
                                temp_dir.glob("**/dashboards/**/*.yaml")
                            ) + list(temp_dir.glob("**/dashboards/*.yaml"))
                            dash_files = list(set(dash_files))
                            logger.reason(
                                f"Patching cross-filters for {len(dash_files)} dashboards"
                            )
                            # Gather all source UUID-to-ID mappings from the archive first
                            source_id_to_uuid_map = (
                                self._extract_chart_uuids_from_archive(temp_dir)
                            )
                            for dash_file in dash_files:
                                self._patch_dashboard_metadata(
                                    dash_file, target_env_id, source_id_to_uuid_map
                                )
                        else:
                            logger.explore(
                                "Cross-filter patching requested but mapping service or target_env_id is missing"
                            )
                    # 3. Re-package
                    logger.reason(
                        f"Re-packaging transformed archive (strip_databases={strip_databases})"
                    )
                    with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
                        for root, dirs, files in os.walk(temp_dir):
                            rel_root = Path(root).relative_to(temp_dir)
                            for file in files:
                                file_path = Path(root) / file
                                if "databases" in rel_root.parts and (
                                    strip_databases
                                    or file_path not in referenced_database_files
                                ):
                                    continue
                                arcname = file_path.relative_to(temp_dir)
                                zf.write(file_path, arcname)
                    logger.reflect("ZIP transformation completed successfully")
                    return True
                except Exception as e:
                    logger.explore(f"Error transforming ZIP: {e}")
                    return False
    # #endregion Core.MigrationEngine.TransformZip

    # #region Core.MigrationEngine.CollectRequiredDatabaseUuidsFromZip [C:2] [TYPE Function] [SEMANTICS migration,archive,database]
    # @BRIEF Inspect an export ZIP and return dataset-referenced source database UUIDs.
    # @POST Returns a set of UUID strings (possibly empty) without mutating the archive.
    def collect_required_database_uuids_from_zip(self, zip_path: str) -> set[str]:
        with belief_scope("MigrationEngine.collect_required_database_uuids_from_zip"):
            with tempfile.TemporaryDirectory() as temp_dir_str:
                temp_dir = Path(temp_dir_str)
                try:
                    with zipfile.ZipFile(zip_path, "r") as zf:
                        zf.extractall(temp_dir)
                    dataset_files = list(temp_dir.glob("**/datasets/**/*.yaml")) + list(
                        temp_dir.glob("**/datasets/*.yaml")
                    )
                    return self._collect_dataset_database_uuids(list(set(dataset_files)))
                except Exception as exc:
                    logger.explore(
                        "Could not inspect archive for database UUIDs",
                        error=str(exc),
                    )
                    return set()
    # #endregion Core.MigrationEngine.CollectRequiredDatabaseUuidsFromZip

    # #region Core.MigrationEngine.ListDatabaseResourcesFromZip [C:2] [TYPE Function] [SEMANTICS migration,archive,database]
    # @BRIEF Return [{uuid, name, file}] for database YAMLs referenced by datasets in a ZIP.
    def list_database_resources_from_zip(self, zip_path: str) -> list[dict[str, str]]:
        with belief_scope("MigrationEngine.list_database_resources_from_zip"):
            resources: list[dict[str, str]] = []
            with tempfile.TemporaryDirectory() as temp_dir_str:
                temp_dir = Path(temp_dir_str)
                try:
                    with zipfile.ZipFile(zip_path, "r") as zf:
                        zf.extractall(temp_dir)
                    dataset_files = list(set(
                        list(temp_dir.glob("**/datasets/**/*.yaml"))
                        + list(temp_dir.glob("**/datasets/*.yaml"))
                    ))
                    required = self._collect_dataset_database_uuids(dataset_files)
                    db_files = list(set(
                        list(temp_dir.glob("**/databases/**/*.yaml"))
                        + list(temp_dir.glob("**/databases/*.yaml"))
                    ))
                    for db_file in db_files:
                        uuid = self._database_yaml_uuid(db_file)
                        if not uuid or uuid not in required:
                            continue
                        name = db_file.stem
                        try:
                            with open(db_file) as stream:
                                data = yaml.safe_load(stream) or {}
                            if isinstance(data.get("database_name"), str) and data["database_name"]:
                                name = data["database_name"]
                        except Exception:
                            pass
                        resources.append({"uuid": uuid, "name": name, "file": db_file.name})
                    # UUIDs referenced by datasets but missing database YAML
                    known = {r["uuid"] for r in resources}
                    for uuid in sorted(required - known):
                        resources.append({"uuid": uuid, "name": uuid[:8] + "…", "file": ""})
                except Exception as exc:
                    logger.explore(
                        "Could not list database resources from archive",
                        error=str(exc),
                    )
            return resources
    # #endregion Core.MigrationEngine.ListDatabaseResourcesFromZip

    # #region Core.MigrationEngine.ReadDatasetContractsFromZip [C:3] [TYPE Function] [SEMANTICS migration,archive,dataset,contract]
    # @ingroup Core
    # @BRIEF Read dataset contracts from an export ZIP in-memory without mutating the archive.
    # @PRE zip_path points to a readable Superset export ZIP.
    # @POST Returns a list of {uuid, database_uuid, catalog, schema, table_name} dicts; skips
    #       unparseable or empty dataset YAMLs and entries without a truthy uuid.
    # @SIDE_EFFECT Reads the ZIP fully into memory; never writes or mutates the archive.
    # @RELATION DEPENDS_ON -> [Core.MigrationEngine]
    def read_dataset_contracts_from_zip(self, zip_path: str) -> list[dict]:
        with belief_scope("MigrationEngine.read_dataset_contracts_from_zip"):
            try:
                contracts: list[dict] = []
                with zipfile.ZipFile(zip_path, "r") as zf:
                    for name in sorted(zf.namelist()):
                        if not name.endswith(".yaml"):
                            continue
                        if "datasets" not in name.split("/"):
                            continue
                        try:
                            data = yaml.safe_load(zf.read(name)) or {}
                        except Exception:
                            continue
                        if not isinstance(data, dict) or not data.get("uuid"):
                            continue
                        contracts.append(
                            {
                                "uuid": data.get("uuid"),
                                "database_uuid": data.get("database_uuid"),
                                "catalog": data.get("catalog"),
                                "schema": data.get("schema"),
                                "table_name": data.get("table_name"),
                            }
                        )
                logger.reflect(
                    "Dataset contracts read from archive",
                    payload={"count": len(contracts), "zip_path": zip_path},
                )
                return contracts
            except Exception as exc:
                logger.explore(
                    "Could not read dataset contracts from archive",
                    payload={"zip_path": zip_path},
                    error=str(exc),
                )
                return []
    # #endregion Core.MigrationEngine.ReadDatasetContractsFromZip

    # #region Core.MigrationEngine.CollectDatasetDatabaseUuids [C:2] [TYPE Function] [SEMANTICS migration,archive,database,mapping]
    # @BRIEF Read source database UUIDs referenced by the datasets in an export archive.
    # @PRE Each path is a readable Superset dataset YAML file.
    # @POST Returns exactly the non-empty database_uuid values declared by dataset YAMLs.
    # @INVARIANT Only declared dataset dependencies can cause a database resource to be imported.
    def _collect_dataset_database_uuids(self, dataset_files: list[Path]) -> set[str]:
        database_uuids: set[str] = set()
        for dataset_file in dataset_files:
            with open(dataset_file) as stream:
                data = yaml.safe_load(stream) or {}
            database_uuid = data.get("database_uuid")
            if isinstance(database_uuid, str) and database_uuid:
                database_uuids.add(database_uuid)
        return database_uuids
    # #endregion Core.MigrationEngine.CollectDatasetDatabaseUuids

    # #region Core.MigrationEngine.DatabaseYamlUuid [C:1] [TYPE Function] [SEMANTICS migration,archive,database]
    # @BRIEF Return the UUID declared by an exported database resource.
    # @POST Returns None for empty or malformed database YAML without making it importable.
    def _database_yaml_uuid(self, file_path: Path) -> str | None:
        with open(file_path) as stream:
            data = yaml.safe_load(stream) or {}
        database_uuid = data.get("uuid")
        return database_uuid if isinstance(database_uuid, str) and database_uuid else None
    # #endregion Core.MigrationEngine.DatabaseYamlUuid

    # #region Core.MigrationEngine.TransformYaml [TYPE Function]
    # @PURPOSE: Replaces database_uuid in a single YAML file.
    # @PARAM  file_path (Path) - Path to the YAML file.
    # @PARAM  db_mapping (Dict[str, str]) - UUID mapping dictionary.
    # @PRE  file_path exists, is readable YAML, and db_mapping contains source->target UUID pairs.
    # @POST  database_uuid is replaced in-place only when source UUID is present in db_mapping.
    # @SIDE_EFFECT  Reads and conditionally rewrites YAML file on disk.
    # @DATA_CONTRACT  Input[(Path file_path, Dict[str,str] db_mapping)] -> Output[None]
    def _transform_yaml(self, file_path: Path, db_mapping: dict[str, str]):
        with belief_scope("MigrationEngine._transform_yaml"):
            if not file_path.exists():
                logger.explore(f"YAML file not found: {file_path}")
                raise FileNotFoundError(str(file_path))
            with open(file_path) as f:
                data = yaml.safe_load(f)
            if not data:
                return
            source_uuid = data.get("database_uuid")
            if source_uuid in db_mapping:
                logger.reason(f"Replacing database UUID in {file_path.name}")
                data["database_uuid"] = db_mapping[source_uuid]
                with open(file_path, "w") as f:
                    yaml.dump(data, f)
                logger.reflect(f"Database UUID patched in {file_path.name}")
    # #endregion Core.MigrationEngine.TransformYaml
    # #region Core.MigrationEngine.TransformDatabaseYaml [TYPE Function]
    # @PURPOSE: Replaces the top-level uuid field in a database YAML file with the target UUID.
    # @PARAM  file_path (Path) - Path to the database YAML file.
    # @PARAM  db_mapping (Dict[str, str]) - UUID mapping dictionary.
    # @PRE  file_path exists, is readable YAML, and db_mapping contains source->target UUID pairs.
    # @POST  uuid is replaced in-place. Superset import_database() will find the target UUID
    #        as an existing DB, skip creation, and populate database_ids for dataset import.
    # @SIDE_EFFECT  Reads and conditionally rewrites YAML file on disk.
    # @RATIONALE Superset matches databases by UUID during import. If the database YAML in the
    #   archive has a UUID that already exists in the target instance, import_database() returns
    #   the existing DB immediately without password validation. Replacing source UUIDs with
    #   target UUIDs avoids both: password errors AND the cascading GENERIC_COMMAND_ERROR
    #   that occurs when databases/ is stripped entirely (empty database_ids blocks datasets).
    def _transform_database_yaml(self, file_path: Path, db_mapping: dict[str, str]):
        with belief_scope("MigrationEngine._transform_database_yaml"):
            if not file_path.exists():
                logger.explore(f"Database YAML file not found: {file_path}")
                raise FileNotFoundError(str(file_path))
            with open(file_path) as f:
                data = yaml.safe_load(f)
            if not data:
                return
            source_uuid = data.get("uuid")
            if source_uuid is None:
                logger.explore(f"Database YAML has no uuid field: {file_path.name}")
                return
            if source_uuid in db_mapping:
                logger.reason(f"Replacing database UUID in {file_path.name}")
                data["uuid"] = db_mapping[source_uuid]
                with open(file_path, "w") as f:
                    yaml.dump(data, f)
                logger.reflect(f"Database UUID patched: {source_uuid} → {db_mapping[source_uuid]} in {file_path.name}")
            else:
                logger.reason(f"Database UUID {source_uuid} not in mapping — keeping as-is in {file_path.name}")
    # #endregion Core.MigrationEngine.TransformDatabaseYaml
    # #region Core.MigrationEngine.ExtractChartUuidsFromArchive [TYPE Function]
    # @PURPOSE: Scans extracted chart YAML files and builds a source chart ID to UUID lookup map.
    # @PRE  temp_dir exists and points to extracted archive root with optional chart YAML resources.
    # @POST  Returns a best-effort Dict[int, str] containing only parseable chart id/uuid pairs.
    # @SIDE_EFFECT  Reads chart YAML files from filesystem; suppresses per-file parsing failures.
    # @DATA_CONTRACT  Input[Path] -> Output[Dict[int,str]]
    # @PARAM  temp_dir (Path) - Root dir of unpacked archive.
    # @RETURN  Dict[int, str] - Mapping of source Integer ID to UUID.
    def _extract_chart_uuids_from_archive(self, temp_dir: Path) -> dict[int, str]:
        with belief_scope("MigrationEngine._extract_chart_uuids_from_archive"):
            # Implementation Note: This is a placeholder for the logic that extracts
            # actual Source IDs. In a real scenario, this involves parsing chart YAMLs
            # or manifesting the export metadata structure where source IDs are stored.
            # For simplicity in US1 MVP, we assume it's read from chart files if present.
            mapping = {}
            chart_files = list(temp_dir.glob("**/charts/**/*.yaml")) + list(
                temp_dir.glob("**/charts/*.yaml")
            )
            for cf in set(chart_files):
                try:
                    with open(cf) as f:
                        cdata = yaml.safe_load(f)
                        if cdata and "id" in cdata and "uuid" in cdata:
                            mapping[cdata["id"]] = cdata["uuid"]
                except Exception:
                    logger.debug("Could not parse chart YAML file %s", cf)
            return mapping
    # #endregion Core.MigrationEngine.ExtractChartUuidsFromArchive
    # #region Core.MigrationEngine.PatchDashboardMetadata [TYPE Function]
    # @PURPOSE: Rewrites dashboard json_metadata chart/dataset integer identifiers using target environment mappings.
    # @PRE  file_path points to dashboard YAML with json_metadata; target_env_id is non-empty; source_map contains source id->uuid.
    # @POST  json_metadata is re-serialized with mapped integer IDs when remote mappings are available; otherwise file remains unchanged.
    # @SIDE_EFFECT  Reads/writes YAML file, performs mapping lookup via mapping_service, emits logs for recoverable/terminal failures.
    # @DATA_CONTRACT  Input[(Path file_path, str target_env_id, Dict[int,str] source_map)] -> Output[None]
    # @PARAM  file_path (Path)
    # @PARAM  target_env_id (str)
    # @PARAM  source_map (Dict[int, str])
    def _patch_dashboard_metadata(
        self, file_path: Path, target_env_id: str, source_map: dict[int, str]
    ):
        with belief_scope("MigrationEngine._patch_dashboard_metadata"):
            try:
                if not file_path.exists():
                    return
                with open(file_path) as f:
                    data = yaml.safe_load(f)
                if not data or "json_metadata" not in data:
                    return
                metadata_str = data["json_metadata"]
                if not metadata_str:
                    return
                # Fetch target UUIDs for everything we know:
                uuids_needed = list(source_map.values())
                logger.reason(
                    f"Resolving {len(uuids_needed)} remote IDs for dashboard metadata patching"
                )
                target_ids = self.mapping_service.get_remote_ids_batch(
                    target_env_id, ResourceType.CHART, uuids_needed
                )
                if not target_ids:
                    logger.reflect(
                        "No remote target IDs found in mapping database for this dashboard."
                    )
                    return
                # Map Source Int -> Target Int
                source_to_target = {}
                missing_targets = []
                for s_id, s_uuid in source_map.items():
                    if s_uuid in target_ids:
                        source_to_target[s_id] = target_ids[s_uuid]
                    else:
                        missing_targets.append(s_id)
                if missing_targets:
                    logger.explore(
                        f"Missing target IDs for source IDs: {missing_targets}. Cross-filters might break."
                    )
                if not source_to_target:
                    logger.reflect("No source IDs matched remotely. Skipping patch.")
                    return
                logger.reason(
                    f"Patching {len(source_to_target)} ID references in json_metadata"
                )
                new_metadata_str = metadata_str
                for s_id, t_id in source_to_target.items():
                    new_metadata_str = re.sub(
                        r'("datasetId"\s*:\s*)' + str(s_id) + r"(\b)",
                        r"\g<1>" + str(t_id) + r"\g<2>",
                        new_metadata_str,
                    )
                    new_metadata_str = re.sub(
                        r'("chartId"\s*:\s*)' + str(s_id) + r"(\b)",
                        r"\g<1>" + str(t_id) + r"\g<2>",
                        new_metadata_str,
                    )
                # Re-parse to validate valid JSON
                data["json_metadata"] = json.dumps(json.loads(new_metadata_str))
                with open(file_path, "w") as f:
                    yaml.dump(data, f)
                logger.reflect(
                    f"Dashboard metadata patched and saved: {file_path.name}"
                )
            except Exception as e:
                logger.explore(f"Metadata patch failed for {file_path.name}: {e}")
    # #endregion Core.MigrationEngine.PatchDashboardMetadata
# #endregion Core.MigrationEngine
# #endregion Core.MigrationEngine.MigrationEngineModule