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()