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.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:
    match = re.search(
        pattern,
        COMPOSE.read_text(encoding="utf-8"),
        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):
    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}"
    )
    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(database_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()