ss-tools/backend/tests/integration/test_dashboard_migration_db_replacement.py

# #region Test.Migration.DatabaseReplacementE2E [C:5] [TYPE Module] [SEMANTICS test,migration,superset,testcontainers,database]
# @BRIEF Verify a real Superset dashboard export is transformed and imported against a mapped target database.
# @RELATION BINDS_TO -> [Core.MigrationEngine]
# @RELATION BINDS_TO -> [Core.DashboardsCrud.SupersetClientImportDashboard]
# @TEST_CONTRACT Dashboard export ZIP + source-to-target database UUID mapping -> imported dataset bound to target database.
# @TEST_FIXTURE Real Superset 4.1.2 and PostgreSQL 16 -> Testcontainers fixtures in integration/conftest.py.
# @TEST_EDGE missing_mapping -> covered by TestMigrationEngine.test_transform_zip_rejects_partial_database_mapping.
# @TEST_EDGE unrelated_database_resource -> covered by TestMigrationEngine.test_transform_zip_keeps_only_referenced_mapped_database.
# @TEST_EDGE identity_mapping -> covered by TestMigrationEngine.test_transform_zip_allows_empty_mapping_as_identity.
# @TEST_INVARIANT MigrationPlugin.strict_db_isolation -> VERIFIED_BY: test_dashboard_import_replaces_source_database.
# @PRE Docker daemon is available and the shared Superset/PostgreSQL fixtures are healthy.
# @POST The imported dataset references the target database ID and no longer references the source database ID.
# @SIDE_EFFECT Creates and removes Superset database, dataset, chart, and dashboard resources.
# @RATIONALE A real export/import cycle is required because mocked orchestration cannot prove Superset resolves the rewritten database UUID during import.
# @REJECTED Asserting only transform_zip call arguments was rejected — it cannot detect an archive that Superset imports against the source database.
from contextlib import suppress
import json
from pathlib import Path
import pytest
from urllib.parse import urlparse
from uuid import uuid4
import zipfile

from sqlalchemy import text
import yaml

from src.core.migration_engine import MigrationEngine


# #region Test.Migration.DatabaseReplacementE2E.ResourceId [C:1] [TYPE Function]
def _resource_id(response: dict) -> int:
    resource_id = response.get("id")
    if resource_id is None:
        result = response.get("result", {})
        resource_id = result.get("id") if isinstance(result, dict) else None
    if resource_id is None:
        raise AssertionError(f"Superset create response has no resource id: {response}")
    return int(resource_id)
# #endregion Test.Migration.DatabaseReplacementE2E.ResourceId


# #region Test.Migration.DatabaseReplacementE2E.DatabaseUuid [C:1] [TYPE Function]
async def _database_uuid(client, database_id: int) -> str:
    response = await client.get_database(database_id)
    result = response.get("result", response)
    return str(result["uuid"])
# #endregion Test.Migration.DatabaseReplacementE2E.DatabaseUuid


# #region Test.Migration.DatabaseReplacementE2E.IsolateUuids [C:2] [TYPE Function]
def _isolate_resource_uuids(zip_path: Path) -> None:
    """Avoid same-instance UUID collisions while preserving the mapped DB UUID."""
    replacements: dict[str, str] = {}
    with zipfile.ZipFile(zip_path) as archive:
        entries = {name: archive.read(name) for name in archive.namelist()}
    for name, content in entries.items():
        normalized = f"/{name}"
        if not name.endswith((".yaml", ".yml")):
            continue
        if not any(f"/{kind}/" in normalized for kind in ("datasets", "charts", "dashboards")):
            continue
        payload = yaml.safe_load(content) or {}
        resource_uuid = payload.get("uuid")
        if isinstance(resource_uuid, str) and resource_uuid:
            replacements[resource_uuid] = str(uuid4())
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as archive:
        for name, content in entries.items():
            updated = content
            if name.endswith((".yaml", ".yml")):
                text_content = content.decode("utf-8")
                for source_uuid, target_uuid in replacements.items():
                    text_content = text_content.replace(source_uuid, target_uuid)
                updated = text_content.encode("utf-8")
            archive.writestr(name, updated)
# #endregion Test.Migration.DatabaseReplacementE2E.IsolateUuids


# #region Test.Migration.DatabaseReplacementE2E.Test [C:4] [TYPE Function] [SEMANTICS test,migration,superset,database]
# @ingroup Test.Migration.DatabaseReplacementE2E
# @BRIEF Export a real dashboard graph, rewrite its database UUID, import it, and inspect the resulting dataset relation.
@pytest.mark.integration
@pytest.mark.asyncio
async def test_dashboard_import_replaces_source_database(
    superset_client,
    superset_db_url,
    pg_engine,
    tmp_path: Path,
):
    suffix = uuid4().hex[:10]
    table_name = f"migration_db_replace_{suffix}"
    dashboard_title = f"Migration DB Replacement {suffix}"
    source_database_id = None
    target_database_id = None
    dataset_id = None
    imported_dataset_id = None
    chart_id = None
    dashboard_id = None
    imported_dashboard_id = None

    parsed = urlparse(superset_db_url)
    database_uri = (
        f"postgresql://test:test@{parsed.hostname}:{parsed.port}/test_translate"
    )

    with pg_engine.begin() as connection:
        connection.execute(
            text(f'CREATE TABLE "{table_name}" (id INTEGER PRIMARY KEY, label VARCHAR(32))')
        )
        connection.execute(
            text(f'INSERT INTO "{table_name}" (id, label) VALUES (1, \'source-row\')')
        )

    try:
        source_database = await superset_client.create_database(
            database_name=f"Migration Source {suffix}",
            sqlalchemy_uri=database_uri,
        )
        source_database_id = _resource_id(source_database)
        target_database = await superset_client.create_database(
            database_name=f"Migration Target {suffix}",
            sqlalchemy_uri=database_uri,
        )
        target_database_id = _resource_id(target_database)

        source_database_uuid = await _database_uuid(
            superset_client, source_database_id
        )
        target_database_uuid = await _database_uuid(
            superset_client, target_database_id
        )
        assert source_database_uuid != target_database_uuid

        dataset = await superset_client.create_dataset(
            table_name=table_name,
            database=source_database_id,
            schema_name="public",
        )
        dataset_id = _resource_id(dataset)

        dashboard = await superset_client.create_dashboard(
            dashboard_title=dashboard_title,
            slug=f"migration-db-replacement-{suffix}",
        )
        dashboard_id = _resource_id(dashboard)

        chart = await superset_client.client.request(
            method="POST",
            endpoint="/chart/",
            data={
                "dashboards": [dashboard_id],
                "datasource_id": dataset_id,
                "datasource_type": "table",
                "slice_name": f"Migration Chart {suffix}",
                "viz_type": "table",
                "params": json.dumps(
                    {
                        "datasource": f"{dataset_id}__table",
                        "viz_type": "table",
                        "all_columns": ["id", "label"],
                        "row_limit": 100,
                    }
                ),
            },
        )
        chart_id = _resource_id(chart)

        await superset_client.client.request(
            method="PUT",
            endpoint=f"/dashboard/{dashboard_id}",
            data={
                "position_json": json.dumps(
                    {
                        f"CHART-{chart_id}": {
                            "id": f"CHART-{chart_id}",
                            "type": "CHART",
                            "meta": {"chartId": chart_id, "width": 12, "height": 50},
                        }
                    }
                )
            },
        )

        export_content, _ = await superset_client.export_dashboard(dashboard_id)
        source_zip = tmp_path / "source-dashboard.zip"
        transformed_zip = tmp_path / "target-dashboard.zip"
        source_zip.write_bytes(export_content)

        transformed = MigrationEngine().transform_zip(
            str(source_zip),
            str(transformed_zip),
            {source_database_uuid: target_database_uuid},
            strip_databases=False,
        )
        assert transformed is True

        # The shared fixture has one Superset metadata DB. Fresh non-database UUIDs
        # emulate a separate target while leaving the production DB rewrite intact.
        _isolate_resource_uuids(transformed_zip)

        await superset_client.import_dashboard(transformed_zip)

        _, all_datasets = await superset_client.get_datasets()
        imported_datasets = [
            item for item in all_datasets if item.get("table_name") == table_name
        ]
        assert len(imported_datasets) == 2
        imported_dataset_ids = [
            int(item["id"])
            for item in imported_datasets
            if int(item["id"]) != dataset_id
        ]
        assert len(imported_dataset_ids) == 1
        imported_dataset_id = imported_dataset_ids[0]
        imported_dataset = await superset_client.get_dataset(imported_dataset_id)
        imported_result = imported_dataset.get("result", imported_dataset)
        imported_database = imported_result.get("database") or {}
        imported_database_id = imported_result.get("database_id") or imported_database.get("id")
        assert int(imported_database_id) == target_database_id
        assert int(imported_database_id) != source_database_id

        _, all_dashboards = await superset_client.get_dashboards()
        imported_dashboards = [
            item
            for item in all_dashboards
            if item.get("dashboard_title") == dashboard_title
        ]
        imported_dashboard_ids = [
            int(item["id"])
            for item in imported_dashboards
            if int(item["id"]) != dashboard_id
        ]
        if imported_dashboard_ids:
            imported_dashboard_id = imported_dashboard_ids[0]
    finally:
        if imported_dashboard_id is not None:
            with suppress(Exception):
                await superset_client.delete_dashboard(imported_dashboard_id)
        if dashboard_id is not None:
            with suppress(Exception):
                await superset_client.delete_dashboard(dashboard_id)
        if chart_id is not None:
            with suppress(Exception):
                await superset_client.delete_chart(chart_id)
        if dataset_id is not None:
            with suppress(Exception):
                await superset_client.delete_dataset(dataset_id)
        if imported_dataset_id is not None:
            with suppress(Exception):
                await superset_client.delete_dataset(imported_dataset_id)
        for database_id in (source_database_id, target_database_id):
            if database_id is not None:
                with suppress(Exception):
                    await superset_client.delete_database(database_id)
        with pg_engine.begin() as connection:
            connection.execute(text(f'DROP TABLE IF EXISTS "{table_name}"'))
# #endregion Test.Migration.DatabaseReplacementE2E.Test

# #endregion Test.Migration.DatabaseReplacementE2E