superset2/tests/integration_tests/versioning/retention_prune_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.
"""Integration coverage for version-history retention pruning.

Exercises ``superset.tasks.version_history_retention`` end-to-end against
a real database: that an aged-out transaction's shadow rows are pruned
while the live row is always preserved, and that the SERIALIZABLE pass
retries on transient serialization failures and gives up after the cap.
"""

from __future__ import annotations

from collections.abc import Iterator
from typing import Any

import pytest
from sqlalchemy_continuum import version_class

from superset.extensions import db
from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
from superset.versioning.changes.table import version_changes_table
from tests.integration_tests.base_tests import SupersetTestCase
from tests.integration_tests.fixtures.birth_names_dashboard import (  # noqa: F401
    load_birth_names_dashboard_with_slices,
    load_birth_names_data,
)


def _get_version_rows(dashboard: Dashboard) -> list[Any]:
    ver_cls = version_class(Dashboard)
    return (
        db.session.query(ver_cls)
        .filter(ver_cls.id == dashboard.id)
        .order_by(ver_cls.transaction_id.asc())
        .all()
    )


def _persist_fixture_state() -> None:
    """Force the fixture's pending INSERTs to commit in their own transaction.

    The birth_names fixture stages charts and the dashboard via session.add()
    but does not commit. Without this, the test's first commit batches the
    INSERTs and UPDATEs into the same Continuum transaction, causing the
    existing version row to be updated in place instead of a new one being
    created.
    """
    db.session.commit()


class TestDashboardVersionRetention(SupersetTestCase):
    """Retention pruning drops shadow rows older than
    ``SUPERSET_VERSION_HISTORY_RETENTION_DAYS`` while preserving live rows."""

    @pytest.fixture(autouse=True)
    def _load_data(
        self,
        load_birth_names_dashboard_with_slices: None,  # noqa: F811
    ) -> Iterator[None]:
        """Establish the capture state this integration class requires.

        Continuum stores its option and SQLAlchemy event listeners globally.
        Initialization unit tests intentionally exercise the kill-switch that
        detaches those listeners, so a mixed or reordered pytest invocation can
        otherwise enter these tests with capture disabled. Reasserting the
        integration suite's prerequisite here makes each test order-independent.
        """
        import sqlalchemy as sa
        from sqlalchemy_continuum import versioning_manager

        from superset.initialization import SupersetAppInitializer

        previous_versioning: bool = bool(
            versioning_manager.options.get("versioning", True)
        )
        listeners_were_attached: bool = sa.event.contains(
            sa.orm.Mapper,
            "after_insert",
            versioning_manager.track_inserts,
        )

        versioning_manager.options["versioning"] = True
        SupersetAppInitializer._add_continuum_write_listeners()
        try:
            yield
        finally:
            if listeners_were_attached:
                SupersetAppInitializer._add_continuum_write_listeners()
            else:
                SupersetAppInitializer._remove_continuum_write_listeners()
            versioning_manager.options["versioning"] = previous_versioning

    def test_retention_prunes_old_rows(self) -> None:
        """``prune_old_versions`` removes shadow rows whose owning
        ``version_transaction.issued_at`` is older than the retention
        window, while preserving every transaction that anchors a live row."""
        from datetime import datetime, timedelta, timezone

        import sqlalchemy as sa

        from superset.extensions import db as _db
        from superset.tasks.version_history_retention import (
            _prune_old_versions_impl,
        )

        _persist_fixture_state()
        dashboard: Dashboard = (
            db.session.query(Dashboard)
            .filter(Dashboard.dashboard_title == "USA Births Names")
            .first()
        )
        assert dashboard is not None

        original_title = dashboard.dashboard_title

        try:
            # Force a few saves so we have ≥ 2 closed shadow rows plus
            # a baseline plus the live row.
            for i in range(3):
                dashboard.dashboard_title = f"USA Births Names retention test {i}"
                db.session.commit()

            rows_before = _get_version_rows(dashboard)
            assert len(rows_before) >= 3, "Expected at least 3 version rows"

            def _version_changes_count() -> int:
                return (
                    _db.session.execute(
                        sa.text("SELECT COUNT(*) FROM version_changes")
                    ).scalar()
                    or 0
                )

            changes_before = _version_changes_count()
            assert changes_before >= 1, (
                "expected version_changes rows from the edits above"
            )

            # Backdate every version_transaction row by 100 days so the
            # prune sees them as old. Skip baseline+live rows; the prune
            # itself preserves them.
            from sqlalchemy_continuum import versioning_manager

            tx_table = versioning_manager.transaction_cls.__table__
            with _db.engine.begin() as conn:
                conn.execute(
                    sa.update(tx_table).values(
                        issued_at=datetime.now(timezone.utc).replace(tzinfo=None)
                        - timedelta(days=100)
                    )
                )

            stats = _prune_old_versions_impl(retention_days=30)
            assert stats.get("pruned_transactions", 0) >= 1, stats

            # version_changes rows for pruned transactions must be gone too.
            # The prune deletes version_transaction rows and relies on the
            # ON DELETE CASCADE FK from version_changes.transaction_id; assert
            # the cascade actually fired rather than orphaning change records.
            db.session.expire_all()
            assert _version_changes_count() < changes_before, (
                "version_changes rows for pruned transactions were not removed "
                f"(before={changes_before}, after={_version_changes_count()})"
            )

            rows_after = _get_version_rows(dashboard)
            # Live row must still exist (this is the only preservation rule)
            live_rows = [r for r in rows_after if r.end_transaction_id is None]
            assert len(live_rows) >= 1, "Live row must never be pruned"
            # Some rows should have been pruned. Closed historical rows —
            # including the synthetic baseline (operation_type=0) — are
            # subject to retention like everything else.
            assert len(rows_after) < len(rows_before), (
                f"Expected fewer rows after prune; before={len(rows_before)} "
                f"after={len(rows_after)}"
            )

        finally:
            dashboard.dashboard_title = original_title
            db.session.commit()

    def test_retention_preserves_live_child_and_m2m_rows(self) -> None:
        """Regression for the live-row preservation BLOCKER.

        A parent-only edit leaves the dashboard's chart associations (the
        M2M ``dashboard_slices_version`` rows) live but anchored at an
        *older* transaction than the dashboard's current live row. The
        prune must preserve that older transaction too — if it scans only
        parent shadows for live rows, it deletes the still-live
        association and the surviving version silently loses its slices.

        The parent-only ``test_retention_prunes_old_rows`` cannot catch
        this: it asserts on the dashboard parent shadow alone, which is
        exactly the blind spot the bug lived in.
        """
        from datetime import datetime, timedelta, timezone

        import sqlalchemy as sa
        from sqlalchemy_continuum import version_class, versioning_manager

        from superset.extensions import db as _db
        from superset.tasks.version_history_retention import _prune_old_versions_impl

        _persist_fixture_state()
        dashboard: Dashboard = (
            db.session.query(Dashboard)
            .filter(Dashboard.dashboard_title == "USA Births Names")
            .first()
        )
        assert dashboard is not None
        assert dashboard.slices, "fixture dashboard must have slices for this test"
        dashboard_id = dashboard.id
        original_title = dashboard.dashboard_title

        m2m: sa.Table = version_class(Dashboard).__table__.metadata.tables[
            "dashboard_slices_version"
        ]

        def _live_m2m_count() -> int:
            with _db.engine.begin() as conn:
                return conn.execute(
                    sa.select(sa.func.count())
                    .select_from(m2m)
                    .where(m2m.c.dashboard_id == dashboard_id)
                    .where(m2m.c.end_transaction_id.is_(None))
                ).scalar_one()

        try:
            # Baseline capture synthesizes the live M2M association rows at
            # the baseline tx. A parent-only edit then opens a newer parent
            # live row while the association rows stay live at the older tx.
            dashboard.dashboard_title = "USA Births Names (parent-only edit)"
            db.session.commit()

            live_before = _live_m2m_count()
            assert live_before >= 1, (
                "expected live dashboard_slices_version rows before prune"
            )

            # Backdate every transaction so the whole chain is older than
            # the window — including the tx anchoring the live association.
            tx_table = versioning_manager.transaction_cls.__table__
            with _db.engine.begin() as conn:
                conn.execute(
                    sa.update(tx_table).values(
                        issued_at=datetime.now(timezone.utc).replace(tzinfo=None)
                        - timedelta(days=100)
                    )
                )

            _prune_old_versions_impl(retention_days=30)

            # The live association rows must survive: their anchoring tx is
            # preserved because they are live, even though it predates the
            # dashboard's current live parent row.
            assert _live_m2m_count() == live_before, (
                "prune deleted live dashboard_slices_version rows — the "
                "surviving version lost its slices (preservation BLOCKER)"
            )
        finally:
            dashboard.dashboard_title = original_title
            db.session.commit()

    def test_retention_preserves_multi_flush_transaction(self) -> None:
        """Pruning preserves #41940's final multi-flush semantic projection."""
        from datetime import datetime, timedelta, timezone

        import sqlalchemy as sa
        from sqlalchemy_continuum import versioning_manager

        from superset.tasks.version_history_retention import _prune_old_versions_impl

        _persist_fixture_state()
        dashboard: Dashboard = (
            db.session.query(Dashboard)
            .filter(Dashboard.dashboard_title == "USA Births Names")
            .one()
        )
        chart: Slice = dashboard.slices[0]
        dashboard_title = dashboard.dashboard_title
        chart_name = chart.slice_name
        tx_table = versioning_manager.transaction_cls.__table__
        dashboard_version_table = version_class(Dashboard).__table__
        chart_version_table = version_class(Slice).__table__
        boundary = db.session.scalar(sa.select(sa.func.max(tx_table.c.id))) or 0

        try:
            dashboard.dashboard_title = "USA Births Names multi-flush dashboard"
            db.session.flush()
            chart.slice_name = "USA Births Names multi-flush chart"
            db.session.commit()

            change_rows = db.session.execute(
                sa.select(
                    version_changes_table.c.transaction_id,
                    version_changes_table.c.entity_kind,
                )
                .where(version_changes_table.c.transaction_id > boundary)
                .where(
                    sa.or_(
                        sa.and_(
                            version_changes_table.c.entity_kind == "dashboard",
                            version_changes_table.c.entity_id == dashboard.id,
                        ),
                        sa.and_(
                            version_changes_table.c.entity_kind == "chart",
                            version_changes_table.c.entity_id == chart.id,
                        ),
                    )
                )
            ).all()
            assert {row.entity_kind for row in change_rows} == {"dashboard", "chart"}
            transaction_ids: set[int] = {row.transaction_id for row in change_rows}
            assert len(transaction_ids) == 1
            transaction_id: int = transaction_ids.pop()
            assert (
                db.session.scalar(
                    sa.select(dashboard_version_table.c.transaction_id)
                    .where(dashboard_version_table.c.id == dashboard.id)
                    .where(dashboard_version_table.c.end_transaction_id.is_(None))
                )
                == transaction_id
            )
            assert (
                db.session.scalar(
                    sa.select(chart_version_table.c.transaction_id)
                    .where(chart_version_table.c.id == chart.id)
                    .where(chart_version_table.c.end_transaction_id.is_(None))
                )
                == transaction_id
            )

            with db.engine.begin() as conn:
                conn.execute(
                    sa.update(tx_table).values(
                        issued_at=datetime.now(timezone.utc).replace(tzinfo=None)
                        - timedelta(days=100)
                    )
                )

            _prune_old_versions_impl(retention_days=30)
            db.session.rollback()

            assert (
                db.session.scalar(
                    sa.select(sa.func.count())
                    .select_from(tx_table)
                    .where(tx_table.c.id == transaction_id)
                )
                == 1
            )
            surviving_kinds = set(
                db.session.scalars(
                    sa.select(version_changes_table.c.entity_kind).where(
                        version_changes_table.c.transaction_id == transaction_id
                    )
                )
            )
            assert {"dashboard", "chart"}.issubset(surviving_kinds)
            assert (
                db.session.scalar(
                    sa.select(sa.func.count())
                    .select_from(dashboard_version_table)
                    .where(dashboard_version_table.c.id == dashboard.id)
                    .where(dashboard_version_table.c.end_transaction_id.is_(None))
                )
                == 1
            )
            assert (
                db.session.scalar(
                    sa.select(sa.func.count())
                    .select_from(chart_version_table)
                    .where(chart_version_table.c.id == chart.id)
                    .where(chart_version_table.c.end_transaction_id.is_(None))
                )
                == 1
            )
        finally:
            dashboard.dashboard_title = dashboard_title
            chart.slice_name = chart_name
            db.session.commit()

    def test_retention_retries_on_serialization_failure(self) -> None:
        """A transient ``OperationalError`` from the SERIALIZABLE pass
        triggers an inline retry; the prune completes on the second
        attempt and the stats dict records the retry count."""
        from datetime import datetime, timedelta, timezone
        from unittest.mock import patch

        import sqlalchemy as sa
        from sqlalchemy.exc import OperationalError

        from superset.tasks import version_history_retention
        from superset.tasks.version_history_retention import (
            _prune_old_versions_impl,
        )

        # Backdate transactions so the prune has work to do.
        _persist_fixture_state()
        dashboard: Dashboard = (
            db.session.query(Dashboard)
            .filter(Dashboard.dashboard_title == "USA Births Names")
            .first()
        )
        original_title = dashboard.dashboard_title
        try:
            for i in range(3):
                dashboard.dashboard_title = f"USA Births Names retry test {i}"
                db.session.commit()

            from sqlalchemy_continuum import versioning_manager

            tx_table = versioning_manager.transaction_cls.__table__
            from superset.extensions import db as _db

            with _db.engine.begin() as conn:
                conn.execute(
                    sa.update(tx_table).values(
                        issued_at=datetime.now(timezone.utc).replace(tzinfo=None)
                        - timedelta(days=100)
                    )
                )

            original_run = version_history_retention._run_prune_pass
            calls: list[int] = []

            def flaky_run(*args: Any, **kwargs: Any) -> dict[str, Any]:
                calls.append(1)
                if len(calls) == 1:
                    raise OperationalError(
                        "SELECT 1", {}, Exception("could not serialize access")
                    )
                return original_run(*args, **kwargs)

            with patch.object(
                version_history_retention, "_run_prune_pass", side_effect=flaky_run
            ):
                # Patch sleep so the test doesn't actually wait through
                # the backoff.
                with patch.object(version_history_retention.time, "sleep"):
                    stats = _prune_old_versions_impl(retention_days=30)

            # At least 2 passes: the first conflicts (call 1) and is
            # retried (call 2). There may be more when the backlog spans
            # multiple id-ordered windows (the prune batches by
            # _MAX_PRUNE_BATCH), but exactly one conflict was injected, so
            # the retry counter must read 1 regardless of batch count.
            assert len(calls) >= 2, (
                f"Expected >= 2 _run_prune_pass calls (>= 1 failure + retry), "
                f"got {len(calls)}"
            )
            assert stats.get("retried") == 1, stats
            assert stats.get("pruned_transactions", 0) >= 1, stats
        finally:
            dashboard.dashboard_title = original_title
            db.session.commit()

    def test_retention_gives_up_after_max_attempts(self) -> None:
        """When every attempt hits ``OperationalError``, the function
        re-raises after the retry cap so the outer Celery wrapper logs
        + returns ``{"error": 1}``."""
        from unittest.mock import patch

        from sqlalchemy.exc import OperationalError

        from superset.tasks import version_history_retention
        from superset.tasks.version_history_retention import (
            _MAX_RETRY_ATTEMPTS,
            _prune_old_versions_impl,
        )

        def always_fail(*args: Any, **kwargs: Any) -> dict[str, Any]:
            raise OperationalError(
                "SELECT 1", {}, Exception("could not serialize access")
            )

        call_count: int = 0

        def counting_fail(*args: Any, **kwargs: Any) -> dict[str, Any]:
            nonlocal call_count
            call_count += 1
            return always_fail(*args, **kwargs)

        with patch.object(
            version_history_retention,
            "_run_prune_pass",
            side_effect=counting_fail,
        ):
            with patch.object(version_history_retention.time, "sleep"):
                with pytest.raises(OperationalError):
                    _prune_old_versions_impl(retention_days=30)

        assert call_count == _MAX_RETRY_ATTEMPTS, (
            f"Expected exactly {_MAX_RETRY_ATTEMPTS} attempts; got {call_count}"
        )