import asyncio from pathlib import Path import pytest from alembic import command from alembic.config import Config from sqlalchemy import text from sqlalchemy.ext.asyncio import create_async_engine def get_alembic_config(database_url: str) -> Config: root_dir = Path(__file__).resolve().parent.parent ini_path = root_dir / "alembic.ini" alembic_cfg = Config(str(ini_path)) # Pass via attributes to avoid '%' escaping issues with passwords alembic_cfg.attributes["sqlalchemy_url"] = database_url return alembic_cfg @pytest.mark.asyncio async def test_alembic_migration_upgrade_and_downgrade_lifecycle(temp_postgres_db: str): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Run upgrade to head await asyncio.to_thread(command.upgrade, alembic_cfg, "head") # 2. Verify alembic_version has the current head and tables exist async with engine.connect() as conn: result = await conn.execute(text("SELECT version_num FROM alembic_version;")) row = result.fetchone() assert row is not None assert row[0] == "0010_dotb_headers" table_check = await conn.execute( text( "SELECT EXISTS (" " SELECT FROM information_schema.tables " " WHERE table_schema = 'public' AND table_name = 'scoring_runs'" ");" ) ) assert table_check.scalar() is True audit_table_check = await conn.execute( text( "SELECT EXISTS (" " SELECT FROM information_schema.tables " " WHERE table_schema = 'public' AND table_name = 'review_audits'" ");" ) ) assert audit_table_check.scalar() is True # Check new columns exist cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'scoring_runs';" ) ) cols = {r[0] for r in cols_result.fetchall()} assert { "ai_score", "final_score", "scoring_mode", "degraded", "ai_metadata", "version", "effective_score", "review_decision", "reviewed_by", "reviewed_at", "review_note", "parent_run_id", "redacted_at", "execution_trace", "dotb_app", "dotb_feature", }.issubset(cols) audit_cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'review_audits';" ) ) audit_cols = {r[0] for r in audit_cols_result.fetchall()} assert { "child_run_id", "effective_score", "redacted_at", }.issubset(audit_cols) outbox_cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'outbox_events';" ) ) outbox_cols = {r[0] for r in outbox_cols_result.fetchall()} assert { "delivery_status", "attempt_count", "next_attempt_at", "last_attempt_at", "delivered_at", "lease_expires_at", "claim_token", "last_error_code", "last_error_detail", "last_http_status", }.issubset(outbox_cols) # 3. Run downgrade to base await asyncio.to_thread(command.downgrade, alembic_cfg, "base") # 4. Verify alembic_version is empty and scoring_runs table is dropped async with engine.connect() as conn: result_after_downgrade = await conn.execute( text("SELECT version_num FROM alembic_version;") ) row_after = result_after_downgrade.fetchone() assert row_after is None table_check_after = await conn.execute( text( "SELECT EXISTS (" " SELECT FROM information_schema.tables " " WHERE table_schema = 'public' AND table_name = 'scoring_runs'" ");" ) ) assert table_check_after.scalar() is False await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_repeated_upgrade_idempotence(temp_postgres_db: str): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # Initial upgrade await asyncio.to_thread(command.upgrade, alembic_cfg, "head") # Second upgrade to head should be a no-op and succeed cleanly await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: result = await conn.execute(text("SELECT version_num FROM alembic_version;")) row = result.fetchone() assert row is not None assert row[0] == "0010_dotb_headers" # Downgrade back to base to keep disposable container clean await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_downgrade_with_populated_review_and_outbox_data( temp_postgres_db: str, ): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Upgrade to head (0005) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") run_id = "00000000-0000-0000-0000-000000000001" terminal_event_id = "00000000-0000-0000-0000-000000000002" rev1_event_id = "00000000-0000-0000-0000-000000000003" rev2_event_id = "00000000-0000-0000-0000-000000000004" audit1_id = "00000000-0000-0000-0000-000000000005" audit2_id = "00000000-0000-0000-0000-000000000006" # 2. Insert populated data representing multiple reviews and terminal outbox event async with engine.begin() as conn: await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " ai_score, final_score, effective_score, scoring_mode, " " degraded, review_decision, reviewed_by, reviewed_at, review_note, " " created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'run-idem-1', " " 'hash-1', 'completed', '{}'::jsonb, 3, 80, " " NULL, 80, 95, 'rules_only', " " false, 'override', 'admin_user', now(), 'Second override note', " " now(), now()" ")" ), {"id": run_id}, ) # Terminal event await conn.execute( text( "INSERT INTO outbox_events (" " id, event_type, event_version, run_id, payload, occurred_at" ") VALUES (" " :id, 'lead_scoring.run.terminal', '1.0.0', :run_id, " " '{}'::jsonb, now()" ")" ), {"id": terminal_event_id, "run_id": run_id}, ) # Two review events for the same run_id (would violate unique without partial index) await conn.execute( text( "INSERT INTO outbox_events (" " id, event_type, event_version, run_id, payload, occurred_at" ") VALUES (" " :id, 'lead_scoring.run.reviewed', '1.0.0', :run_id, " " '{\"version\": 2}'::jsonb, now()" ")" ), {"id": rev1_event_id, "run_id": run_id}, ) await conn.execute( text( "INSERT INTO outbox_events (" " id, event_type, event_version, run_id, payload, occurred_at" ") VALUES (" " :id, 'lead_scoring.run.reviewed', '1.0.0', :run_id, " " '{\"version\": 3}'::jsonb, now()" ")" ), {"id": rev2_event_id, "run_id": run_id}, ) # Review audits await conn.execute( text( "INSERT INTO review_audits (" " id, run_id, tenant_id, decision, actor_id, actor_role, idempotency_key, " " payload_hash, previous_version, new_version, calculated_score, effective_score, " " note, created_at" ") VALUES (" " :id, :run_id, 'tenant_1', 'override', 'admin_user', 'administrator', 'idem-1', " " 'hash1', 1, 2, 80, 90, 'first note', now()" ")" ), {"id": audit1_id, "run_id": run_id}, ) await conn.execute( text( "INSERT INTO review_audits (" " id, run_id, tenant_id, decision, actor_id, actor_role, idempotency_key, " " payload_hash, previous_version, new_version, calculated_score, effective_score, " " note, created_at" ") VALUES (" " :id, :run_id, 'tenant_1', 'override', 'admin_user', 'administrator', 'idem-2', " " 'hash2', 2, 3, 80, 95, 'second note', now()" ")" ), {"id": audit2_id, "run_id": run_id}, ) # 3. Downgrade to 0004_medium_ai_pipeline (prior revision) await asyncio.to_thread(command.downgrade, alembic_cfg, "0004_medium_ai_pipeline") # 4. Verify post-downgrade state async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0004_medium_ai_pipeline" # review_audits table should be gone audit_check = await conn.execute( text( "SELECT EXISTS (" " SELECT FROM information_schema.tables " " WHERE table_schema = 'public' AND table_name = 'review_audits'" ");" ) ) assert audit_check.scalar() is False # scoring_runs should have kept existing pre-0005 data while columns were dropped run_res = await conn.execute( text("SELECT id, status, final_score FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) run_row = run_res.fetchone() assert run_row is not None assert run_row[0] == run_id assert run_row[1] == "completed" assert run_row[2] == 80 # outbox_events should retain terminal event while review events are purged events_res = await conn.execute( text("SELECT id, event_type FROM outbox_events WHERE run_id = :id;"), {"id": run_id} ) events = events_res.fetchall() assert len(events) == 1 assert events[0][0] == terminal_event_id assert events[0][1] == "lead_scoring.run.terminal" # 5. Upgrade back to head to verify round-trip await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0010_dotb_headers" # Downgrade to base to clean up await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_0006_downgrade_and_upgrade_with_reevaluate_data( temp_postgres_db: str, ): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Upgrade to head (0006) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") parent_run_id = "00000000-0000-0000-0000-000000000010" child_run_id = "00000000-0000-0000-0000-000000000011" audit_id = "00000000-0000-0000-0000-000000000012" async with engine.begin() as conn: # Insert parent run (superseded) await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " ai_score, final_score, effective_score, scoring_mode, " " degraded, review_decision, reviewed_by, reviewed_at, review_note, " " parent_run_id, created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'parent-idem-1', " " 'hash-p', 'superseded', '{}'::jsonb, 2, 60, " " NULL, 60, NULL, 'rules_only', " " false, 're_evaluate', 'admin_user', now(), 'Re-evaluating with new data', " " NULL, now(), now()" ")" ), {"id": parent_run_id}, ) # Insert child run with parent_run_id await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " ai_score, final_score, effective_score, scoring_mode, " " degraded, review_decision, reviewed_by, reviewed_at, review_note, " " parent_run_id, created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'child-idem-1', " " 'hash-c', 'pending', '{}'::jsonb, 1, NULL, " " NULL, NULL, NULL, NULL, " " NULL, NULL, NULL, NULL, NULL, " " :parent_id, now(), now()" ")" ), {"id": child_run_id, "parent_id": parent_run_id}, ) # Insert review audit with NULL effective_score and child_run_id await conn.execute( text( "INSERT INTO review_audits (" " id, run_id, tenant_id, decision, actor_id, actor_role, idempotency_key, " " payload_hash, previous_version, new_version, " " calculated_score, effective_score, " " child_run_id, note, created_at" ") VALUES (" " :id, :parent_id, 'tenant_1', 're_evaluate', 'admin_user', " " 'administrator', 'rev-idem-1', " " 'hash-rev', 1, 2, 60, NULL, " " :child_id, 'Re-evaluating with new data', now()" ")" ), {"id": audit_id, "parent_id": parent_run_id, "child_id": child_run_id}, ) # 2. Downgrade to 0005_review_api_and_audit await asyncio.to_thread(command.downgrade, alembic_cfg, "0005_review_api_and_audit") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0005_review_api_and_audit" # Check that audit effective_score defaulted to 0 on downgrade (not NULL) audit_res = await conn.execute( text("SELECT effective_score FROM review_audits WHERE id = :id;"), {"id": audit_id}, ) assert audit_res.scalar() == 0 # Check that parent_run_id column no longer exists cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'scoring_runs';" ) ) cols = {r[0] for r in cols_result.fetchall()} assert "parent_run_id" not in cols # 3. Upgrade back to head (0009) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0010_dotb_headers" # Downgrade to base await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_0007_downgrade_and_upgrade_with_outbox_data( temp_postgres_db: str, ): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Upgrade to head (0007) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") run_id = "00000000-0000-0000-0000-000000000020" event_id = "00000000-0000-0000-0000-000000000021" async with engine.begin() as conn: await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'run-idem-20', " " 'hash-20', 'completed', '{}'::jsonb, 1, 85, " " now(), now()" ")" ), {"id": run_id}, ) await conn.execute( text( "INSERT INTO outbox_events (" " id, event_type, event_version, run_id, payload, occurred_at, " " delivery_status, attempt_count, next_attempt_at, claim_token, " " last_error_code, last_error_detail, last_http_status" ") VALUES (" " :id, 'lead_scoring.run.terminal', '1.0.0', :run_id, '{}'::jsonb, now(), " " 'dead_letter', 5, now(), 'token-20', " " 'retries_exhausted', 'Max retry limit reached', 500" ")" ), {"id": event_id, "run_id": run_id}, ) # 2. Downgrade to 0006_reevaluate_child_lineage await asyncio.to_thread(command.downgrade, alembic_cfg, "0006_reevaluate_child_lineage") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0006_reevaluate_child_lineage" # Check outbox_events retained original core columns and record outbox_res = await conn.execute( text("SELECT id, event_type, run_id FROM outbox_events WHERE id = :id;"), {"id": event_id}, ) row = outbox_res.fetchone() assert row is not None assert row[0] == event_id assert row[1] == "lead_scoring.run.terminal" assert row[2] == run_id # Check that 0007 columns are dropped cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'outbox_events';" ) ) cols = {r[0] for r in cols_result.fetchall()} assert "delivery_status" not in cols assert "attempt_count" not in cols assert "next_attempt_at" not in cols # 3. Upgrade back to head (0009) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0010_dotb_headers" # Check default values populated on upgraded existing row outbox_res = await conn.execute( text("SELECT id, delivery_status, attempt_count FROM outbox_events WHERE id = :id;"), {"id": event_id}, ) row = outbox_res.fetchone() assert row is not None assert row[1] == "pending" assert row[2] == 0 # Downgrade to base await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_0008_downgrade_and_upgrade_with_retention_data( temp_postgres_db: str, ): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Upgrade to head (0008) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") run_id = "00000000-0000-0000-0000-000000000030" audit_id = "00000000-0000-0000-0000-000000000031" async with engine.begin() as conn: await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " redacted_at, created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'run-idem-30', " " 'hash-30', 'completed', '{}'::jsonb, 1, 85, " " now(), now(), now()" ")" ), {"id": run_id}, ) await conn.execute( text( "INSERT INTO review_audits (" " id, run_id, tenant_id, decision, actor_id, actor_role, idempotency_key, " " payload_hash, previous_version, new_version, calculated_score, effective_score, " " redacted_at, created_at" ") VALUES (" " :id, :run_id, 'tenant_1', 'approve', 'redacted', 'agent', 'rev-idem-30', " " 'redacted', 1, 2, 85, 85, now(), now()" ")" ), {"id": audit_id, "run_id": run_id}, ) # 2. Downgrade to 0007_outbox_delivery_state await asyncio.to_thread(command.downgrade, alembic_cfg, "0007_outbox_delivery_state") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0007_outbox_delivery_state" # Check that redacted_at columns are dropped from scoring_runs and review_audits run_cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'scoring_runs';" ) ) run_cols = {r[0] for r in run_cols_result.fetchall()} assert "redacted_at" not in run_cols audit_cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'review_audits';" ) ) audit_cols = {r[0] for r in audit_cols_result.fetchall()} assert "redacted_at" not in audit_cols # Ensure base records are still preserved run_res = await conn.execute( text("SELECT id, status FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) assert run_res.fetchone() is not None # 3. Upgrade back to head (0009) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0010_dotb_headers" # Check redacted_at is now NULL for existing rows run_check = await conn.execute( text("SELECT id, redacted_at FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) row = run_check.fetchone() assert row is not None assert row[1] is None # Downgrade to base await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_0009_downgrade_and_upgrade_with_execution_trace( temp_postgres_db: str, ): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Upgrade to head (0009) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") run_id = "00000000-0000-0000-0000-000000000040" async with engine.begin() as conn: await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " execution_trace, created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'run-idem-40', " " 'hash-40', 'completed', '{}'::jsonb, 1, 85, " ' \'{"status": "completed", "calculation": {"rule_score": 85}}\'::jsonb, ' " now(), now()" ")" ), {"id": run_id}, ) # 2. Downgrade to 0008_retention_redaction_state await asyncio.to_thread(command.downgrade, alembic_cfg, "0008_retention_redaction_state") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0008_retention_redaction_state" # Check that execution_trace column is dropped run_cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'scoring_runs';" ) ) run_cols = {r[0] for r in run_cols_result.fetchall()} assert "execution_trace" not in run_cols # Ensure base records are still preserved run_res = await conn.execute( text("SELECT id, status FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) assert run_res.fetchone() is not None # 3. Upgrade back to head (0010) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0010_dotb_headers" # Check execution_trace is now NULL for existing rows run_check = await conn.execute( text("SELECT id, execution_trace FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) row = run_check.fetchone() assert row is not None assert row[1] is None # Downgrade to base await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose() @pytest.mark.asyncio async def test_alembic_migration_0010_downgrade_and_upgrade_with_dotb_headers( temp_postgres_db: str, ): alembic_cfg = get_alembic_config(temp_postgres_db) engine = create_async_engine(temp_postgres_db) # 1. Upgrade to head (0010) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") run_id = "00000000-0000-0000-0000-000000000050" async with engine.begin() as conn: await conn.execute( text( "INSERT INTO scoring_runs (" " id, tenant_id, subject_ref, scoring_config_version, idempotency_key, " " payload_hash, status, snapshot_data, version, rule_score, " " dotb_app, dotb_feature, created_at, updated_at" ") VALUES (" " :id, 'tenant_1', 'lead_1', '1.0.0', 'run-idem-50', " " 'hash-50', 'completed', '{}'::jsonb, 1, 85, " " 'dotb_omni', 'lead-scoring', " " now(), now()" ")" ), {"id": run_id}, ) # 2. Downgrade to 0009_execution_trace await asyncio.to_thread(command.downgrade, alembic_cfg, "0009_execution_trace") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0009_execution_trace" # Check that dotb_app and dotb_feature columns are dropped run_cols_result = await conn.execute( text( "SELECT column_name FROM information_schema.columns " "WHERE table_schema = 'public' AND table_name = 'scoring_runs';" ) ) run_cols = {r[0] for r in run_cols_result.fetchall()} assert "dotb_app" not in run_cols assert "dotb_feature" not in run_cols # Ensure base records are still preserved run_res = await conn.execute( text("SELECT id, status FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) assert run_res.fetchone() is not None # 3. Upgrade back to head (0010) await asyncio.to_thread(command.upgrade, alembic_cfg, "head") async with engine.connect() as conn: res = await conn.execute(text("SELECT version_num FROM alembic_version;")) assert res.scalar() == "0010_dotb_headers" # Check columns are restored and NULL for existing row run_check = await conn.execute( text("SELECT id, dotb_app, dotb_feature FROM scoring_runs WHERE id = :id;"), {"id": run_id}, ) row = run_check.fetchone() assert row is not None assert row[1] is None assert row[2] is None # Downgrade to base await asyncio.to_thread(command.downgrade, alembic_cfg, "base") await engine.dispose()