from __future__ import annotations import os import subprocess import uuid from pathlib import Path import psycopg2 import pytest from sqlalchemy import create_engine, inspect from sqlalchemy.engine import make_url ROOT = Path(__file__).resolve().parents[1] EXPECTED_BASELINE_TABLES = { "data_orders", "data_products", "metadata_review_records", "metadata_version_history", "task_list", "users", } EXPECTED_UPGRADED_TABLES = { "datasource_credentials", "datasource_credential_audit_events", "workflow_schedules", "workflow_runs", "workflow_task_runs", "workflow_engine_bindings", "workflow_plan_audits", "runner_task_executions", "workflow_canary_evidence", "workflow_gateway_operations", "workflow_migration_states", "workflow_dual_runs", "workflow_reconciliation_reports", "workflow_cutover_operations", "rule_sql_staging_receipts", "ingestion_sources", "source_artifacts", "ingestion_jobs", "evidence_fragments", "extraction_candidates", "data_elements", "data_element_versions", "candidate_decisions", "ontologies", "ontology_versions", "ontology_domain_links", "ontology_change_sets", "ontology_publish_runs", "governance_responsibility_scopes", "governance_responsibility_assignments", "governance_responsibility_audit_events", "catalog_snapshots", } def test_data_research_ingestion_migration_is_additive_and_constrained(): migration = ( ROOT / "migrations" / "versions" / "20260722_100_data_research_ingestion.py" ).read_text(encoding="utf-8") assert 'revision = "20260722_100"' in migration assert 'down_revision = "20260720_100"' in migration for table in ( "ingestion_sources", "source_artifacts", "ingestion_jobs", "evidence_fragments", "extraction_candidates", ): assert f"CREATE TABLE IF NOT EXISTS public.{table}" in migration assert "idempotency_key" in migration assert "parser_version" in migration assert "locator JSONB" in migration assert "DROP TABLE" not in migration.upper() def test_catalog_execution_migration_adds_attempts_and_snapshots(): migration = ( ROOT / "migrations" / "versions" / "20260729_280_catalog_ingestion_execution.py" ).read_text(encoding="utf-8") assert 'revision = "20260729_280"' in migration assert 'down_revision = "20260729_270"' in migration assert "ADD COLUMN attempt_count" in migration assert "ADD COLUMN failure_stage" in migration assert "CREATE TABLE public.catalog_snapshots" in migration assert "UNIQUE (job_uid, attempt)" in migration assert "DROP TABLE" not in migration.upper() def test_data_element_migration_adds_versioned_governance_tables(): migration = ( ROOT / "migrations" / "versions" / "20260722_105_data_research_elements.py" ).read_text(encoding="utf-8") assert 'revision = "20260722_105"' in migration assert 'down_revision = "20260722_100"' in migration for table in ("data_elements", "data_element_versions", "candidate_decisions"): assert f"CREATE TABLE IF NOT EXISTS public.{table}" in migration assert "UNIQUE (data_element_uid, version)" in migration assert "DROP TABLE" not in migration.upper() def test_ontology_migration_adds_versioned_control_plane_tables(): migration = ( ROOT / "migrations" / "versions" / "20260722_110_data_research_ontology.py" ).read_text(encoding="utf-8") assert 'revision = "20260722_110"' in migration assert 'down_revision = "20260722_105"' in migration for table in ( "ontologies", "ontology_versions", "ontology_domain_links", "ontology_change_sets", "ontology_publish_runs", ): assert f"CREATE TABLE IF NOT EXISTS public.{table}" in migration assert "UNIQUE (ontology_uid, version)" in migration assert "DROP TABLE" not in migration.upper() def test_alembic_configuration_is_environment_only(): ini = (ROOT / "alembic.ini").read_text(encoding="utf-8") env = (ROOT / "migrations" / "env.py").read_text(encoding="utf-8") assert "sqlalchemy.url" not in ini assert "SQLALCHEMY_DATABASE_URI" in env assert "DATABASE_URL" in env assert "password" not in env.lower() def test_baseline_migration_is_non_destructive(): baseline = ( ROOT / "migrations" / "versions" / "20260716_01_baseline.py" ).read_text(encoding="utf-8") assert "def upgrade" in baseline assert "def downgrade" in baseline assert "drop_table" not in baseline for table in EXPECTED_BASELINE_TABLES: assert table in baseline @pytest.mark.integration def test_alembic_upgrade_is_repeatable_and_downgrade_preserves_tables(): admin_url = os.environ.get("TEST_POSTGRES_ADMIN_URL") if not admin_url: pytest.skip("TEST_POSTGRES_ADMIN_URL is not configured") parsed = make_url(admin_url) database_name = f"dataops_migration_{uuid.uuid4().hex[:12]}" target_url = parsed.set(database=database_name).render_as_string( hide_password=False ) connection = psycopg2.connect(admin_url) connection.autocommit = True try: with connection.cursor() as cursor: cursor.execute(f'CREATE DATABASE "{database_name}"') env = os.environ.copy() env["SQLALCHEMY_DATABASE_URI"] = target_url command = [str(ROOT / ".venv" / "bin" / "alembic"), "-c", "alembic.ini"] subprocess.run(command + ["upgrade", "head"], cwd=ROOT, env=env, check=True) subprocess.run(command + ["upgrade", "head"], cwd=ROOT, env=env, check=True) engine = create_engine(target_url) try: tables = set(inspect(engine).get_table_names(schema="public")) assert EXPECTED_BASELINE_TABLES | {"alembic_version"} <= tables assert tables >= EXPECTED_UPGRADED_TABLES finally: engine.dispose() subprocess.run(command + ["downgrade", "-1"], cwd=ROOT, env=env, check=True) engine = create_engine(target_url) try: tables = set(inspect(engine).get_table_names(schema="public")) assert tables >= EXPECTED_BASELINE_TABLES finally: engine.dispose() finally: with connection.cursor() as cursor: cursor.execute( "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " "WHERE datname = %s AND pid <> pg_backend_pid()", (database_name,), ) cursor.execute(f'DROP DATABASE IF EXISTS "{database_name}"') connection.close() @pytest.mark.integration @pytest.mark.parametrize( ("legacy_rows", "diagnostic"), ( ((1, 1), "duplicate rule_run_id"), ((101,), "rows above 100"), ), ) def test_upgrade_from_150_rejects_malformed_legacy_samples( legacy_rows, diagnostic, ): admin_url = os.environ.get("TEST_POSTGRES_ADMIN_URL") if not admin_url: pytest.skip("TEST_POSTGRES_ADMIN_URL is not configured") parsed = make_url(admin_url) database_name = f"dataops_evidence_{uuid.uuid4().hex[:12]}" target_url = parsed.set(database=database_name).render_as_string( hide_password=False ) admin = psycopg2.connect(admin_url) admin.autocommit = True try: with admin.cursor() as cursor: cursor.execute(f'CREATE DATABASE "{database_name}"') env = os.environ.copy() env["SQLALCHEMY_DATABASE_URI"] = target_url command = [ str(ROOT / ".venv" / "bin" / "alembic"), "-c", "alembic.ini", ] subprocess.run( command + ["upgrade", "20260723_150"], cwd=ROOT, env=env, check=True, ) database = psycopg2.connect(target_url) try: database.autocommit = True with database.cursor() as cursor: cursor.execute("SET session_replication_role = replica") run_id = str(uuid.uuid4()) cursor.execute( """ INSERT INTO public.rule_runs ( id, deployment_id, component_binding_id, rule_version_id, plan_hash, status, correlation_id ) VALUES (%s, %s, %s, %s, %s, 'failed', %s) """, ( run_id, str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4()), "a" * 64, str(uuid.uuid4()), ), ) for count in legacy_rows: cursor.execute( """ INSERT INTO public.rule_violation_samples ( id, rule_run_id, artifact_ref, sample_count, redaction_policy, expires_at ) VALUES (%s, %s, %s, %s, 'legacy-v1', NOW()) """, ( str(uuid.uuid4()), run_id, f"minio://legacy/{uuid.uuid4()}", count, ), ) cursor.execute("SET session_replication_role = origin") finally: database.close() failed = subprocess.run( command + ["upgrade", "head"], cwd=ROOT, env=env, text=True, capture_output=True, ) assert failed.returncode != 0 assert diagnostic in (failed.stdout + failed.stderr) database = psycopg2.connect(target_url) try: with database.cursor() as cursor: cursor.execute( "SELECT version_num FROM public.alembic_version" ) assert cursor.fetchone()[0] == "20260723_150" finally: database.close() finally: with admin.cursor() as cursor: cursor.execute( "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " "WHERE datname = %s AND pid <> pg_backend_pid()", (database_name,), ) cursor.execute(f'DROP DATABASE IF EXISTS "{database_name}"') admin.close() @pytest.mark.integration def test_concurrent_alembic_upgrades_are_serialized(): admin_url = os.environ.get("TEST_POSTGRES_ADMIN_URL") if not admin_url: pytest.skip("TEST_POSTGRES_ADMIN_URL is not configured") parsed = make_url(admin_url) database_name = f"dataops_concurrent_{uuid.uuid4().hex[:12]}" target_url = parsed.set(database=database_name).render_as_string( hide_password=False ) connection = psycopg2.connect(admin_url) connection.autocommit = True try: with connection.cursor() as cursor: cursor.execute(f'CREATE DATABASE "{database_name}"') env = os.environ.copy() env["SQLALCHEMY_DATABASE_URI"] = target_url command = [ str(ROOT / ".venv" / "bin" / "alembic"), "-c", "alembic.ini", "upgrade", "head", ] processes = [ subprocess.Popen( command, cwd=ROOT, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, ) for _ in range(2) ] results = [process.communicate(timeout=120) for process in processes] failures = [ output for process, (output, _) in zip(processes, results) if process.returncode != 0 ] assert failures == [] engine = create_engine(target_url) try: assert set(inspect(engine).get_table_names(schema="public")) >= ( EXPECTED_UPGRADED_TABLES | {"alembic_version"} ) finally: engine.dispose() finally: with connection.cursor() as cursor: cursor.execute( "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " "WHERE datname = %s AND pid <> pg_backend_pid()", (database_name,), ) cursor.execute(f'DROP DATABASE IF EXISTS "{database_name}"') connection.close()