"""Managed PostgreSQL qualification for benchmark-run persistence and restart behavior. This test is enabled only by the disposable PostgreSQL boundary in managed-validation.yml. It creates one uniquely named database on that boundary, then drops only that database if it still exists and is still owned by the expected validation role. """ from __future__ import annotations import os import threading import uuid from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime from pathlib import Path import pytest import sqlalchemy as sa from alembic.config import Config from sqlalchemy import create_engine, func, inspect, select, text from sqlalchemy.engine import URL, make_url from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from alembic import command from modelforge_api.domain.benchmarking import BenchmarkRunStartRequest from modelforge_api.persistence.models import ( ArtifactSet, ArtifactSetMember, BenchmarkRun, Deployment, GatewayRequest, LabExecutionRun, ) from modelforge_api.services.benchmark_runs import BenchmarkRunService from modelforge_api.services.serving import ServingService from modelforge_api.services.transient_payloads import MemoryPayloadStore from tests.test_dispatch_context import operator_fixture _PRIOR = "20260924_0038" _CURRENT = "20260924_0039" _LIFECYCLE_OPT_IN = "managed-validation-v1" def _managed_test_urls() -> tuple[URL, URL, URL]: """Fail closed unless the dedicated disposable PostgreSQL boundary opted in.""" if os.getenv("MODELFORGE_TEST_POSTGRES_DATABASE_LIFECYCLE") is None: pytest.skip("requires the managed-validation disposable PostgreSQL lifecycle opt-in") if os.getenv("MODELFORGE_TEST_POSTGRES_DATABASE_LIFECYCLE") != _LIFECYCLE_OPT_IN: pytest.fail("managed PostgreSQL database lifecycle opt-in has an unexpected value") raw_urls = ( os.getenv("MODELFORGE_TEST_POSTGRES_URL"), os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL"), os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL"), ) if not all(raw_urls): pytest.fail("managed PostgreSQL lifecycle opt-in requires all three test role URLs") urls = tuple(make_url(value) for value in raw_urls if value is not None) expected_users = ("postgres", "modelforge", "modelforge_runtime") expected_passwords = ( os.getenv("MODELFORGE_POSTGRES_ADMIN_PASSWORD"), os.getenv("MODELFORGE_MIGRATION_DB_PASSWORD"), os.getenv("MODELFORGE_RUNTIME_DB_PASSWORD"), ) labels = ("admin", "migration owner", "runtime") if not all(expected_passwords): pytest.fail("managed PostgreSQL role passwords are unavailable in this test process") for url, username, password, label in zip( urls, expected_users, expected_passwords, labels, strict=True ): if ( url.drivername != "postgresql+psycopg" or url.host != "postgres" or url.port != 5432 or url.database != "modelforge" or url.username != username or url.query ): pytest.fail( f"managed PostgreSQL {label} URL does not match the isolated CI test boundary" ) if url.password != password: pytest.fail( f"managed PostgreSQL {label} URL does not match its dedicated test role password" ) return urls # type: ignore[return-value] def _current_revision(engine) -> str: with engine.connect() as connection: return connection.scalar(text("SELECT version_num FROM alembic_version")) def _alembic_config(database_url: URL) -> Config: config = Config() config.set_main_option("script_location", str(Path(__file__).resolve().parents[1] / "alembic")) rendered_url = database_url.render_as_string(hide_password=False).replace("%", "%%") config.set_main_option("sqlalchemy.url", rendered_url) return config def _insert_legacy_run( connection: sa.Connection, *, run_id: uuid.UUID, deployment_id: uuid.UUID, suite_id: uuid.UUID, ) -> None: """Seed the exact pre-0039 row shape to prove the additive migration preserves history.""" legacy_runs = sa.table( "benchmark_runs", sa.column("id", sa.Uuid()), sa.column("deployment_id", sa.Uuid()), sa.column("suite_id", sa.Uuid()), sa.column("model_revision_sha", sa.String(64)), sa.column("artifact_sha256", sa.String(64)), sa.column("runtime_type", sa.String(32)), sa.column("runtime_version", sa.String(128)), sa.column("runtime_image_digest", sa.String(71)), sa.column("launch_arguments", sa.JSON()), sa.column("suite_revision", sa.String(128)), sa.column("dataset_revision", sa.String(255)), sa.column("environment_fingerprint", sa.JSON()), sa.column("environment_digest", sa.String(64)), sa.column("seed", sa.Integer()), sa.column("results", sa.JSON()), sa.column("status", sa.String(32)), sa.column("comparable", sa.Boolean()), sa.column("stale_reason", sa.Text()), sa.column("started_at", sa.DateTime(timezone=True)), sa.column("completed_at", sa.DateTime(timezone=True)), ) now = datetime.now(UTC) connection.execute( legacy_runs.insert().values( id=run_id, deployment_id=deployment_id, suite_id=suite_id, model_revision_sha="a" * 40, artifact_sha256="b" * 64, runtime_type="onnx", runtime_version="1.0", runtime_image_digest=None, launch_arguments={}, suite_revision="1", dataset_revision="builtin:legacy-benchmark-fixture-v1", environment_fingerprint={}, environment_digest="c" * 64, seed=None, results={"legacy_fixture": True}, status="completed", comparable=True, stale_reason=None, started_at=now, completed_at=now, ) ) def test_managed_postgresql_benchmark_migration_race_and_restart( monkeypatch: pytest.MonkeyPatch, ) -> None: """Exercise migration 0039 and concurrent/restarted starts against real PostgreSQL.""" admin_url, owner_base_url, runtime_base_url = _managed_test_urls() database_name = f"test_benchmark_run_{uuid.uuid4().hex}" admin_engine = create_engine(admin_url, isolation_level="AUTOCOMMIT", pool_pre_ping=True) created_database = False owner_engine = None runtime_engine = None try: with admin_engine.connect() as connection: already_exists = connection.scalar( text("SELECT 1 FROM pg_database WHERE datname = :database_name"), {"database_name": database_name}, ) assert already_exists is None connection.execute(text(f'CREATE DATABASE "{database_name}" OWNER modelforge')) created_database = True owner_url = owner_base_url.set(database=database_name) runtime_url = runtime_base_url.set(database=database_name) owner_engine = create_engine(owner_url, pool_pre_ping=True) runtime_engine = create_engine( runtime_url, pool_pre_ping=True, connect_args={ "connect_timeout": 15, "options": "-c statement_timeout=30000 -c lock_timeout=15000", }, ) migration = _alembic_config(owner_url) command.upgrade(migration, _PRIOR) legacy_run_id = uuid.uuid4() with Session(owner_engine) as setup_session: serving, target_request, selected = operator_fixture(setup_session) benchmark_service = BenchmarkRunService(serving) advertised_suite = benchmark_service.suites()[0] suite = benchmark_service._suite(advertised_suite.id) request = BenchmarkRunStartRequest( capability="rag.embedding@1", suite_id=suite.id, target=target_request.target, deployment_id=selected.capability_deployment_id, experiment_route_id=selected.id, ) artifact_set = setup_session.get(ArtifactSet, selected.deployment.artifact_set_id) artifact_id = setup_session.scalar( select(ArtifactSetMember.artifact_id) .where(ArtifactSetMember.artifact_set_id == selected.deployment.artifact_set_id) .limit(1) ) assert artifact_set is not None and artifact_id is not None legacy_deployment = Deployment( name=f"benchmark-migration-fixture-{uuid.uuid4().hex}", revision_id=artifact_set.revision_id, artifact_id=artifact_id, derived_artifact_id=None, runtime_profile_id=selected.deployment.runtime_profile_id, compute_node_id=selected.deployment.compute_node_id, channel="candidate", status="draft", capability_contract_ids=[], ) setup_session.add(legacy_deployment) setup_session.flush() legacy_deployment_id = legacy_deployment.id setup_session.commit() deployment_settings = serving._deployment_settings manifests = serving.manifests with owner_engine.begin() as connection: _insert_legacy_run( connection, run_id=legacy_run_id, deployment_id=legacy_deployment_id, suite_id=suite.id, ) command.upgrade(migration, "head") assert _current_revision(owner_engine) == _CURRENT column_metadata = { column["name"]: column for column in inspect(owner_engine).get_columns("benchmark_runs") } upgraded_columns = set(column_metadata) assert column_metadata["deployment_id"]["nullable"] is True assert { "capability_deployment_id", "target_snapshot", "actor_subject", "idempotency_key", "request_fingerprint", "failure_code", "cancellation_requested_at", "created_at", "generation", } <= upgraded_columns constraints = inspect(owner_engine) assert "ck_benchmark_runs_operator_identity" in { constraint["name"] for constraint in constraints.get_check_constraints("benchmark_runs") } assert "uq_benchmark_runs_actor_idempotency" in { constraint["name"] for constraint in constraints.get_unique_constraints("benchmark_runs") } capability_fk = next( foreign_key for foreign_key in constraints.get_foreign_keys("benchmark_runs") if foreign_key["name"] == "fk_benchmark_runs_capability_deployment" ) assert capability_fk["referred_table"] == "capability_deployments" assert capability_fk["options"]["ondelete"] == "RESTRICT" with Session(owner_engine) as session: historical = session.get(BenchmarkRun, legacy_run_id) assert historical is not None assert historical.deployment_id == legacy_deployment_id assert historical.actor_subject is None assert historical.capability_deployment_id is None assert historical.created_at is not None assert historical.generation == 1 with owner_engine.connect() as connection: historical_values = dict( connection.execute( select(BenchmarkRun.__table__).where(BenchmarkRun.id == legacy_run_id) ) .mappings() .one() ) historical_values["id"] = uuid.uuid4() historical_values["actor_subject"] = "partial-identity-must-fail" with pytest.raises(IntegrityError), owner_engine.begin() as connection: connection.execute(BenchmarkRun.__table__.insert().values(historical_values)) with runtime_engine.connect() as connection: assert connection.scalar(text("SELECT version_num FROM alembic_version")) == _CURRENT for privilege in ("SELECT", "INSERT", "UPDATE"): assert connection.scalar( text("SELECT has_table_privilege(current_user, 'benchmark_runs', :privilege)"), {"privilege": privilege}, ) for privilege in ("DELETE", "TRUNCATE"): assert not connection.scalar( text("SELECT has_table_privilege(current_user, 'benchmark_runs', :privilege)"), {"privilege": privilege}, ) original_provenance = BenchmarkRunService._provenance provenance_barrier = threading.Barrier(2) def synchronized_provenance(service, benchmark_request): provenance = original_provenance(service, benchmark_request) provenance_barrier.wait(timeout=30) return provenance monkeypatch.setattr(BenchmarkRunService, "_provenance", synchronized_provenance) def start_once() -> uuid.UUID: with Session(runtime_engine) as session: serving = ServingService( session, deployment_settings, manifests, MemoryPayloadStore(), ) return ( BenchmarkRunService(serving) .start( request, actor_subject="managed-postgres-operator", idempotency_key="same-benchmark-start", ) .id ) with ThreadPoolExecutor(max_workers=2) as pool: run_ids = list(pool.map(lambda _index: start_once(), range(2))) assert run_ids[0] == run_ids[1] benchmark_run_id = run_ids[0] with Session(runtime_engine) as session: benchmark = session.get(BenchmarkRun, benchmark_run_id) assert benchmark is not None assert benchmark.actor_subject == "managed-postgres-operator" assert benchmark.idempotency_key == "same-benchmark-start" assert benchmark.capability_deployment_id == request.deployment_id assert benchmark.target_snapshot["deployment_id"] == str(request.deployment_id) assert ( session.scalar( select(func.count()) .select_from(BenchmarkRun) .where( BenchmarkRun.actor_subject == "managed-postgres-operator", BenchmarkRun.idempotency_key == "same-benchmark-start", ) ) == 1 ) child_id = uuid.UUID(benchmark.results["execution"]["active_lab_run_id"]) child = session.get(LabExecutionRun, child_id) assert child is not None assert child.actor_subject == "managed-postgres-operator" assert child.idempotency_key == f"benchmark-{benchmark_run_id.hex}-00" assert session.get(GatewayRequest, child.gateway_request_id) is not None assert ( session.scalar( select(func.count()) .select_from(LabExecutionRun) .where( LabExecutionRun.actor_subject == "managed-postgres-operator", LabExecutionRun.idempotency_key == f"benchmark-{benchmark_run_id.hex}-00", ) ) == 1 ) # Recreate the engine and service after the race: the committed parent/child link must # remain readable and replaying the same actor/key must not create another child. runtime_engine.dispose() runtime_engine = create_engine( runtime_url, pool_pre_ping=True, connect_args={ "connect_timeout": 15, "options": "-c statement_timeout=30000 -c lock_timeout=15000", }, ) with Session(runtime_engine) as session: serving = ServingService( session, deployment_settings, manifests, MemoryPayloadStore(), ) service = BenchmarkRunService(serving) recovered = service.get(benchmark_run_id, actor_subject="managed-postgres-operator") replay = service.start( request, actor_subject="managed-postgres-operator", idempotency_key="same-benchmark-start", ) assert recovered.id == replay.id == benchmark_run_id assert replay.contract.deployment_id == str(request.deployment_id) assert len(service.list(actor_subject="managed-postgres-operator")) == 1 assert ( session.scalar( select(func.count()) .select_from(LabExecutionRun) .where(LabExecutionRun.actor_subject == "managed-postgres-operator") ) == 1 ) # The downgrade must refuse to erase operator history and leave the migrated schema live. with pytest.raises(RuntimeError, match="Operator benchmark history exists"): command.downgrade(migration, _PRIOR) assert _current_revision(owner_engine) == _CURRENT with Session(owner_engine) as session: assert session.get(BenchmarkRun, benchmark_run_id) is not None assert session.get(BenchmarkRun, legacy_run_id) is not None finally: if runtime_engine is not None: runtime_engine.dispose() if owner_engine is not None: owner_engine.dispose() if created_database: with admin_engine.connect() as connection: database_owner = connection.scalar( text( "SELECT pg_get_userbyid(datdba) FROM pg_database " "WHERE datname = :database_name" ), {"database_name": database_name}, ) if database_owner is not None: if database_owner != "modelforge": raise RuntimeError( "refusing to remove the qualification database after its owner changed" ) connection.execute(text(f'DROP DATABASE "{database_name}" WITH (FORCE)')) admin_engine.dispose()