superset2/tests/integration_tests/migrations/composite_pk_round_trip__tests.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Schema round-trip tests for the composite-PK association-tables migration
(revision 2bee73611e32). Builds the pre-migration shape against an in-memory
SQLite engine, runs the migration's ``upgrade()``, asserts the post-upgrade
shape, runs ``downgrade()``, asserts the prior shape is restored (modulo the
documented FK NOT NULL asymmetry), and re-runs ``upgrade()`` to verify
idempotency.

This is run against an isolated in-memory engine via Alembic's
``MigrationContext`` so the test does not perturb the project's test DB.

Cross-backend (Postgres/MySQL) verification is handled by CI's
test-postgres / test-mysql shards running ``superset db upgrade``. This
file covers the SQLite slice.
"""

from importlib import import_module
from typing import Any

import sqlalchemy as sa
from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import inspect

# Import the migration module under test.
_migration = import_module(
    "superset.migrations.versions."
    "2026-05-01_23-36_2bee73611e32_composite_pk_association_tables"
)
AFFECTED_TABLES = _migration.AFFECTED_TABLES
TABLES_WITH_PRE_EXISTING_UNIQUE = _migration.TABLES_WITH_PRE_EXISTING_UNIQUE
TABLES_WITH_NULLABLE_FKS = _migration.TABLES_WITH_NULLABLE_FKS


def _build_pre_migration_schema(engine: sa.engine.Engine) -> None:
    """Recreate the eight tables in their pre-migration shape (surrogate
    ``id INTEGER PRIMARY KEY`` plus an optional ``UNIQUE(fk1, fk2)`` on the
    two tables that previously carried one). FK columns are NULLABLE on
    the six tables that historically allowed NULLs — fidelity matters:
    with ``nullable=False`` here, the post-upgrade NOT NULL assertions
    pass trivially rather than because the migration promoted anything,
    and the NULL-row cleanup path can't be exercised. FKs to parent
    tables are omitted to keep the test self-contained — we're testing
    schema transformations, not FK enforcement."""
    md = sa.MetaData()
    for t in AFFECTED_TABLES:
        nullable = t.name in TABLES_WITH_NULLABLE_FKS
        cols: list[sa.Column] = [
            sa.Column("id", sa.Integer, primary_key=True),
            sa.Column(t.fk1, sa.Integer, nullable=nullable),
            sa.Column(t.fk2, sa.Integer, nullable=nullable),
        ]
        constraints: list[sa.SchemaItem] = []
        if t.name in TABLES_WITH_PRE_EXISTING_UNIQUE:
            constraints.append(sa.UniqueConstraint(t.fk1, t.fk2))
        sa.Table(t.name, md, *cols, *constraints)
    md.create_all(engine)


def _shape(engine: sa.engine.Engine, table: str) -> dict[str, Any]:
    """Return a structural summary for asserting equality across runs."""
    insp = inspect(engine)
    pk = insp.get_pk_constraint(table).get("constrained_columns", [])
    columns = sorted(c["name"] for c in insp.get_columns(table))
    uniques = sorted(
        tuple(sorted(uc.get("column_names", [])))
        for uc in insp.get_unique_constraints(table)
    )
    return {"columns": columns, "pk": sorted(pk), "uniques": uniques}


def _run_with_alembic_context(engine: sa.engine.Engine, fn) -> None:
    """Run ``fn()`` (the migration's upgrade/downgrade body) inside a fresh
    Alembic ``MigrationContext`` bound to ``engine``. Patches the
    migration module's ``op`` to point at this context so its
    ``op.get_bind()`` and ``op.batch_alter_table`` calls execute against
    the in-memory engine."""
    with engine.connect() as conn:
        ctx = MigrationContext.configure(conn)
        ops = Operations(ctx)
        original_op = _migration.op
        _migration.op = ops  # type: ignore[attr-defined]
        try:
            fn()
        finally:
            _migration.op = original_op  # type: ignore[attr-defined]


def test_round_trip_against_in_memory_sqlite() -> None:
    """Round-trip: pre-migration → upgrade → downgrade → upgrade again.

    Asserts:
    - Post-upgrade shape: no ``id``, composite PK on (fk1, fk2), no
      UNIQUE(fk1, fk2) on the two tables that previously carried one.
    - Post-downgrade shape: ``id`` restored, PK back on (id), UNIQUE
      re-added on the two tables. (FK columns remain NOT NULL — the
      documented intentional asymmetry.)
    - Post-re-upgrade idempotency: shape matches the first post-upgrade.
    """
    engine = sa.create_engine("sqlite:///:memory:")
    _build_pre_migration_schema(engine)

    _run_with_alembic_context(engine, _migration.upgrade)

    for t in AFFECTED_TABLES:
        s = _shape(engine, t.name)
        assert "id" not in s["columns"], f"{t.name}: id still present post-upgrade: {s}"
        assert s["pk"] == sorted([t.fk1, t.fk2]), (
            f"{t.name}: PK is {s['pk']}, expected {sorted([t.fk1, t.fk2])}"
        )
        assert tuple(sorted([t.fk1, t.fk2])) not in s["uniques"], (
            f"{t.name}: redundant UNIQUE not dropped post-upgrade: {s['uniques']}"
        )

    post_upgrade_shape = {t.name: _shape(engine, t.name) for t in AFFECTED_TABLES}

    _run_with_alembic_context(engine, _migration.downgrade)

    for t in AFFECTED_TABLES:
        s = _shape(engine, t.name)
        assert "id" in s["columns"], f"{t.name}: id not restored post-downgrade: {s}"
        assert s["pk"] == ["id"], f"{t.name}: PK is {s['pk']}, expected ['id']"
        if t.name in TABLES_WITH_PRE_EXISTING_UNIQUE:
            assert tuple(sorted([t.fk1, t.fk2])) in s["uniques"], (
                f"{t.name}: UNIQUE not restored post-downgrade: {s['uniques']}"
            )

    _run_with_alembic_context(engine, _migration.upgrade)

    re_upgrade_shape = {t.name: _shape(engine, t.name) for t in AFFECTED_TABLES}
    # Compare with `!=` on the (unhashable) per-table shapes rather than
    # set-differencing `.items()` — the latter raises TypeError while
    # formatting this message and would hide the real regression.
    shape_diff = {
        name: (re_upgrade_shape[name], post_upgrade_shape.get(name))
        for name in re_upgrade_shape
        if re_upgrade_shape[name] != post_upgrade_shape.get(name)
    }
    assert not shape_diff, (
        "Re-upgrade shape differs from initial upgrade shape — "
        f"migration is not idempotent. Per-table diffs (re-upgrade, initial): "
        f"{shape_diff}"
    )


def test_upgrade_scrubs_null_fks_and_duplicates() -> None:
    """The pre-flight data surgery is the migration's riskiest half — and
    it must be deletable-detectable: this test fails if
    ``_delete_null_fk_rows`` or ``_dedupe_by_min_id`` is removed from
    ``upgrade()``.

    Seeds a nullable-FK junction (``slice_user``) with NULL-FK rows and
    duplicate ``(fk1, fk2)`` pairs in the true pre-migration shape, runs
    the upgrade, and asserts exactly the distinct non-NULL pairs survive
    (the composite PK could not even be created otherwise).
    """
    engine = sa.create_engine("sqlite:///:memory:")
    _build_pre_migration_schema(engine)

    md = sa.MetaData()
    slice_user = sa.Table("slice_user", md, autoload_with=engine)
    with engine.begin() as conn:
        conn.execute(
            slice_user.insert(),
            [
                {"id": 1, "user_id": 1, "slice_id": 1},  # keeper (MIN id)
                {"id": 2, "user_id": 1, "slice_id": 1},  # duplicate pair
                {"id": 3, "user_id": 1, "slice_id": 1},  # duplicate pair
                {"id": 4, "user_id": 2, "slice_id": 2},  # distinct keeper
                {"id": 5, "user_id": None, "slice_id": 3},  # NULL fk1
                {"id": 6, "user_id": 3, "slice_id": None},  # NULL fk2
            ],
        )

    _run_with_alembic_context(engine, _migration.upgrade)

    with engine.connect() as conn:
        survivors = sorted(
            conn.execute(sa.text("SELECT user_id, slice_id FROM slice_user")).fetchall()
        )
    assert survivors == [(1, 1), (2, 2)], (
        f"expected the two distinct non-NULL pairs to survive, got {survivors}"
    )


def test_migration_module_constants_are_consistent() -> None:
    """Sanity-check the migration module's exported constants. Catches
    accidental edits that misalign AFFECTED_TABLES with the auxiliary sets."""
    affected_names = {t.name for t in AFFECTED_TABLES}
    assert _migration.TABLES_WITH_PRE_EXISTING_UNIQUE.issubset(affected_names)
    assert _migration.TABLES_WITH_NULLABLE_FKS.issubset(affected_names)
    # Order is alphabetical (deterministic for review/bisection).
    assert [t.name for t in AFFECTED_TABLES] == sorted(affected_names)