DataOps-platform/tests/integration/test_rule_artifact_migration_upgrade.py

from __future__ import annotations

import logging
import os
import re
from pathlib import Path

import polars as pl
import pytest
from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, text
from sqlalchemy.engine import make_url
from sqlalchemy.exc import IntegrityError

from app.core.common.identifiers import new_governance_uid
from tests.runner.test_artifacts import FakeMinio, _store

ROOT = Path(__file__).resolve().parents[2]
COMPOSE = ROOT / "deploy" / "docker" / "docker-compose.yml"


def _compose_value(pattern: str) -> str:
    source = re.sub(
        r"\$\{[A-Za-z_][A-Za-z0-9_]*:-([^}]+)\}",
        r"\1",
        COMPOSE.read_text(encoding="utf-8"),
    )
    match = re.search(
        pattern,
        source,
        flags=re.DOTALL,
    )
    assert match is not None
    return match.group(1)


def _upgrade(database_url: str, revision: str) -> None:
    previous = os.environ.get("DATABASE_URL")
    root_logger = logging.getLogger()
    root_handlers = list(root_logger.handlers)
    root_level = root_logger.level
    logger_disabled = {
        name: logger.disabled
        for name, logger in logging.Logger.manager.loggerDict.items()
        if isinstance(logger, logging.Logger)
    }
    os.environ["DATABASE_URL"] = database_url
    try:
        command.upgrade(Config(str(ROOT / "alembic.ini")), revision)
    finally:
        root_logger.handlers[:] = root_handlers
        root_logger.setLevel(root_level)
        for name, disabled in logger_disabled.items():
            logging.getLogger(name).disabled = disabled
        if previous is None:
            os.environ.pop("DATABASE_URL", None)
        else:
            os.environ["DATABASE_URL"] = previous


def test_old_140_upgrades_to_durable_handoff_and_enforces_cas(
    tmp_path, wp06_postgres_identities
):
    from app.runner.artifacts import PostgresArtifactResolver

    platform_user = _compose_value(
        r"\n  postgres:.*?POSTGRES_USER:\s*([^\s]+)"
    )
    platform_password = _compose_value(
        r"\n  postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
    )
    platform_port = _compose_value(r'"(15432):5432"')
    admin_url = (
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/postgres"
    )
    database_name = f"task5_migration_{new_governance_uid().replace('-', '')}"
    database_url = (
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/{database_name}"
    )
    migration_url = make_url(wp06_postgres_identities.migration_url).set(
        database=database_name
    ).render_as_string(hide_password=False)
    admin = create_engine(admin_url, isolation_level="AUTOCOMMIT")
    engine = None
    try:
        with admin.connect() as connection:
            connection.execute(text(f'CREATE DATABASE "{database_name}"'))
        _upgrade(database_url, "20260723_140")
        engine = create_engine(database_url, pool_pre_ping=True)
        dataflow_version_id = new_governance_uid()
        deployment_id = new_governance_uid()
        schema_id = new_governance_uid()
        binding_id = new_governance_uid()
        binding_hash = "b" * 64
        correlation_id = new_governance_uid()
        old_artifact_id = new_governance_uid()
        old_ref = (
            f"minio://dataops-rules/rules/{correlation_id}/"
            f"{new_governance_uid()}.parquet"
        )
        with engine.begin() as connection:
            columns_before = {
                row[0]
                for row in connection.execute(
                    text(
                        """
                        SELECT column_name
                        FROM information_schema.columns
                        WHERE table_schema = 'public'
                          AND table_name = 'rule_run_artifacts'
                        """
                    )
                )
            }
            assert "handoff_status" not in columns_before
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_versions (
                        id, dataflow_uid, version_no, name, dataflow_spec,
                        input_schema_hashes, output_schema_hash, status
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:uid AS uuid), 1, 'migration',
                        '{}'::jsonb, '[]'::jsonb, :schema_hash, 'released'
                    )
                    """
                ),
                {
                    "id": dataflow_version_id,
                    "uid": new_governance_uid(),
                    "schema_hash": "a" * 64,
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_deployments (
                        id, dataflow_version_id, environment,
                        deployment_config, status
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:version_id AS uuid), 'test',
                        '{}'::jsonb, 'active'
                    )
                    """
                ),
                {"id": deployment_id, "version_id": dataflow_version_id},
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.data_schema_snapshots (
                        id, schema_ref, schema_hash, fields, source_revision
                    ) VALUES (
                        CAST(:id AS uuid), 'migration:id', :schema_hash,
                        CAST(:fields AS jsonb), 'old-140'
                    )
                    """
                ),
                {
                    "id": schema_id,
                    "schema_hash": "a" * 64,
                    "fields": (
                        '[{"name":"id","type":"integer",'
                        '"nullable":false}]'
                    ),
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.dataflow_dataset_bindings (
                        id, dataflow_deployment_id, logical_ref,
                        object_kind, object_ref, schema_snapshot_id, dialect,
                        access_mode, write_mode, binding_hash
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:deployment_id AS uuid),
                        'output', 'parquet_artifact', 'migration-output',
                        CAST(:schema_id AS uuid), 'parquet', 'write',
                        'append', :binding_hash
                    )
                    """
                ),
                {
                    "id": binding_id,
                    "deployment_id": deployment_id,
                    "schema_id": schema_id,
                    "binding_hash": binding_hash,
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_run_artifacts (
                        id, correlation_id, binding_id, artifact_ref,
                        artifact_digest, row_count, schema_hash, schema_fields,
                        artifact_kind, expires_at
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:correlation_id AS uuid),
                        CAST(:binding_id AS uuid), :artifact_ref,
                        :digest, 1, :schema_hash, CAST(:fields AS jsonb),
                        'output', CURRENT_TIMESTAMP + INTERVAL '1 hour'
                    )
                    """
                ),
                {
                    "id": old_artifact_id,
                    "correlation_id": correlation_id,
                    "binding_id": binding_id,
                    "artifact_ref": old_ref,
                    "digest": "d" * 64,
                    "schema_hash": "a" * 64,
                    "fields": (
                        '[{"name":"id","type":"integer",'
                        '"nullable":false}]'
                    ),
                },
            )

        engine.dispose()
        engine = None
        _upgrade(migration_url, "head")
        engine = create_engine(database_url, pool_pre_ping=True)
        with engine.begin() as connection:
            migrated = connection.execute(
                text(
                    """
                    SELECT binding_hash, handoff_status, ready_at
                    FROM public.rule_run_artifacts
                    WHERE id = CAST(:id AS uuid)
                    """
                ),
                {"id": old_artifact_id},
            ).mappings().one()
            assert migrated["binding_hash"] == binding_hash
            assert migrated["handoff_status"] == "ready"
            assert migrated["ready_at"] is not None
            connection.execute(
                text(
                    """
                    DELETE FROM public.rule_run_artifacts
                    WHERE id = CAST(:id AS uuid)
                    """
                ),
                {"id": old_artifact_id},
            )

        store = _store(FakeMinio())
        resolver = PostgresArtifactResolver(engine, store)
        schema_fields = [
            {"name": "id", "type": "integer", "nullable": False}
        ]
        same_path = tmp_path / "same.parquet"
        conflict_path = tmp_path / "conflict.parquet"
        pl.DataFrame({"id": [1]}).write_parquet(same_path)
        pl.DataFrame({"id": [2]}).write_parquet(conflict_path)
        first = resolver.publish_path(
            str(same_path),
            binding_id=binding_id,
            binding_hash=binding_hash,
            correlation_id=correlation_id,
            kind="output",
            ttl_seconds=300,
            schema_fields=schema_fields,
        )
        repeated = resolver.publish_path(
            str(same_path),
            binding_id=binding_id,
            binding_hash=binding_hash,
            correlation_id=correlation_id,
            kind="output",
            ttl_seconds=300,
            schema_fields=schema_fields,
        )
        assert repeated["artifact_ref"] == first["artifact_ref"]
        assert len(store.client.objects) == 1
        with pytest.raises(ValueError, match="immutable|digest"):
            resolver.publish_path(
                str(conflict_path),
                binding_id=binding_id,
                binding_hash=binding_hash,
                correlation_id=correlation_id,
                kind="output",
                ttl_seconds=300,
                schema_fields=schema_fields,
            )
        assert len(store.client.objects) == 1

        with pytest.raises(IntegrityError), engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_run_artifacts (
                        id, correlation_id, binding_id, artifact_ref,
                        artifact_digest, row_count, schema_hash,
                        schema_fields, artifact_kind, binding_hash,
                        handoff_status, expires_at
                    ) VALUES (
                        CAST(:id AS uuid),
                        CAST(:correlation_id AS uuid),
                        CAST(:binding_id AS uuid), :artifact_ref,
                        :artifact_digest, 1, :schema_hash,
                        CAST(:fields AS jsonb), 'output', :binding_hash,
                        'pending',
                        CURRENT_TIMESTAMP + INTERVAL '5 minutes'
                    )
                    """
                ),
                {
                    "id": new_governance_uid(),
                    "correlation_id": correlation_id,
                    "binding_id": binding_id,
                    "artifact_ref": old_ref,
                    "artifact_digest": "e" * 64,
                    "schema_hash": "a" * 64,
                    "fields": (
                        '[{"name":"id","type":"integer",'
                        '"nullable":false}]'
                    ),
                    "binding_hash": binding_hash,
                },
            )
    finally:
        if engine is not None:
            engine.dispose()
        with admin.connect() as connection:
            connection.execute(
                text(
                    """
                    SELECT pg_terminate_backend(pid)
                    FROM pg_stat_activity
                    WHERE datname = :database_name
                      AND pid <> pg_backend_pid()
                    """
                ),
                {"database_name": database_name},
            )
            connection.execute(text(f'DROP DATABASE IF EXISTS "{database_name}"'))
        admin.dispose()


def test_old_180_cleanup_claims_become_expired_and_cas_takeover_ready():
    platform_user = _compose_value(
        r"\n  postgres:.*?POSTGRES_USER:\s*([^\s]+)"
    )
    platform_password = _compose_value(
        r"\n  postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
    )
    platform_port = _compose_value(r'"(15432):5432"')
    admin_url = (
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/postgres"
    )
    database_name = f"task6_claim_{new_governance_uid().replace('-', '')}"
    database_url = (
        f"postgresql+psycopg2://{platform_user}:{platform_password}"
        f"@127.0.0.1:{platform_port}/{database_name}"
    )
    admin = create_engine(admin_url, isolation_level="AUTOCOMMIT")
    engine = None
    try:
        with admin.connect() as connection:
            connection.execute(text(f'CREATE DATABASE "{database_name}"'))
        _upgrade(database_url, "20260723_180")
        engine = create_engine(database_url, pool_pre_ping=True)
        rule_run_id = new_governance_uid()
        sample_id = new_governance_uid()
        receipt_id = new_governance_uid()
        old_sample_claim = new_governance_uid()
        old_receipt_claim = new_governance_uid()
        with engine.begin() as connection:
            expiry_columns = connection.execute(
                text(
                    """
                    SELECT table_name
                    FROM information_schema.columns
                    WHERE table_schema = 'public'
                      AND table_name IN (
                          'rule_violation_samples',
                          'rule_sql_staging_receipts'
                      )
                      AND column_name = 'cleanup_claim_expires_at'
                    """
                )
            ).all()
            assert expiry_columns == []
            connection.execute(text("SET session_replication_role = replica"))
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_runs (
                        id, deployment_id, component_binding_id,
                        rule_version_id, plan_hash, status, correlation_id
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:deployment_id AS uuid),
                        CAST(:component_id AS uuid),
                        CAST(:rule_version_id AS uuid),
                        :plan_hash, 'failed', CAST(:correlation_id AS uuid)
                    )
                    """
                ),
                {
                    "id": rule_run_id,
                    "deployment_id": new_governance_uid(),
                    "component_id": new_governance_uid(),
                    "rule_version_id": new_governance_uid(),
                    "plan_hash": "a" * 64,
                    "correlation_id": new_governance_uid(),
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_violation_samples (
                        id, rule_run_id, artifact_ref, sample_count,
                        redaction_policy, expires_at, cleanup_claim
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:rule_run_id AS uuid),
                        :artifact_ref, 1, 'all-fields',
                        CURRENT_TIMESTAMP - INTERVAL '1 hour',
                        CAST(:cleanup_claim AS uuid)
                    )
                    """
                ),
                {
                    "id": sample_id,
                    "rule_run_id": rule_run_id,
                    "artifact_ref": (
                        f"minio://legacy/{new_governance_uid()}.parquet"
                    ),
                    "cleanup_claim": old_sample_claim,
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_sql_staging_receipts (
                        id, producer_rule_run_id, deployment_id,
                        correlation_id, output_binding_id,
                        output_binding_hash, relation_ref, relation_digest,
                        commit_outcome, status, expires_at, cleanup_claim
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:run_id AS uuid),
                        CAST(:deployment_id AS uuid),
                        CAST(:correlation_id AS uuid),
                        CAST(:binding_id AS uuid), :binding_hash,
                        :relation_ref, :relation_digest,
                        'committed', 'expired',
                        CURRENT_TIMESTAMP - INTERVAL '1 hour',
                        CAST(:cleanup_claim AS uuid)
                    )
                    """
                ),
                {
                    "id": receipt_id,
                    "run_id": rule_run_id,
                    "deployment_id": new_governance_uid(),
                    "correlation_id": new_governance_uid(),
                    "binding_id": new_governance_uid(),
                    "binding_hash": "b" * 64,
                    "relation_ref": "legacy.receipt",
                    "relation_digest": "c" * 64,
                    "cleanup_claim": old_receipt_claim,
                },
            )
            connection.execute(text("SET session_replication_role = origin"))

        engine.dispose()
        engine = None
        _upgrade(database_url, "20260723_190")
        engine = create_engine(database_url, pool_pre_ping=True)
        new_sample_claim = new_governance_uid()
        new_receipt_claim = new_governance_uid()
        with engine.begin() as connection:
            migrated = connection.execute(
                text(
                    """
                    SELECT
                        (
                            SELECT cleanup_claim_expires_at
                            FROM public.rule_violation_samples
                            WHERE id = CAST(:sample_id AS uuid)
                        ) AS sample_claim_expires_at,
                        (
                            SELECT cleanup_claim_expires_at
                            FROM public.rule_sql_staging_receipts
                            WHERE id = CAST(:receipt_id AS uuid)
                        ) AS receipt_claim_expires_at,
                        CURRENT_TIMESTAMP AS observed_at
                    """
                ),
                {"sample_id": sample_id, "receipt_id": receipt_id},
            ).mappings().one()
            assert migrated["sample_claim_expires_at"] is not None
            assert migrated["receipt_claim_expires_at"] is not None
            assert migrated["sample_claim_expires_at"] <= migrated["observed_at"]
            assert migrated["receipt_claim_expires_at"] <= migrated["observed_at"]

            sample_takeover = connection.execute(
                text(
                    """
                    UPDATE public.rule_violation_samples
                    SET cleanup_claim = CAST(:new_claim AS uuid),
                        cleanup_claim_expires_at =
                            CURRENT_TIMESTAMP + INTERVAL '5 minutes'
                    WHERE id = CAST(:id AS uuid)
                      AND cleanup_claim = CAST(:old_claim AS uuid)
                      AND cleanup_claim_expires_at <= CURRENT_TIMESTAMP
                    RETURNING cleanup_claim
                    """
                ),
                {
                    "id": sample_id,
                    "old_claim": old_sample_claim,
                    "new_claim": new_sample_claim,
                },
            ).scalar_one()
            receipt_takeover = connection.execute(
                text(
                    """
                    UPDATE public.rule_sql_staging_receipts
                    SET cleanup_claim = CAST(:new_claim AS uuid),
                        cleanup_claim_expires_at =
                            CURRENT_TIMESTAMP + INTERVAL '5 minutes'
                    WHERE id = CAST(:id AS uuid)
                      AND cleanup_claim = CAST(:old_claim AS uuid)
                      AND cleanup_claim_expires_at <= CURRENT_TIMESTAMP
                    RETURNING cleanup_claim
                    """
                ),
                {
                    "id": receipt_id,
                    "old_claim": old_receipt_claim,
                    "new_claim": new_receipt_claim,
                },
            ).scalar_one()
            assert str(sample_takeover) == new_sample_claim
            assert str(receipt_takeover) == new_receipt_claim
    finally:
        if engine is not None:
            engine.dispose()
        with admin.connect() as connection:
            connection.execute(
                text(
                    """
                    SELECT pg_terminate_backend(pid)
                    FROM pg_stat_activity
                    WHERE datname = :database_name
                      AND pid <> pg_backend_pid()
                    """
                ),
                {"database_name": database_name},
            )
            connection.execute(text(f'DROP DATABASE IF EXISTS "{database_name}"'))
        admin.dispose()