"""Isolated PostgreSQL 0025 upgrade/backup restore and credential race qualification.""" from __future__ import annotations import hashlib import importlib import os import threading import uuid from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime, timedelta from pathlib import Path import pytest from alembic.config import Config from sqlalchemy import create_engine, select, text from sqlalchemy.engine import make_url from sqlalchemy.orm import Session from modelforge_api.domain.serving import ( ServiceClientCreated, ServiceClientResponse, ServiceCredentialRevoke, ServiceCredentialRotate, ) from modelforge_api.persistence.models import AuditEvent, ServiceCredential from modelforge_api.services.audit import AuditWriter from modelforge_api.services.manifest_registry import ManifestRegistry from modelforge_api.services.recovery_fingerprint import ( audit_fingerprint_differences_are_compatible, fingerprint_compatibility, semantic_fingerprint, ) from modelforge_api.services.serving import ServingError, ServingService from modelforge_api.settings import Settings from tests.test_configuration_postgres import _database_tool @pytest.mark.parametrize( "competitors", [("rotate", "rotate"), ("rotate", "revoke"), ("revoke", "revoke")], ) def test_project_upgrade_restore_and_credential_rotation_races( tmp_path: Path, competitors: tuple[str, str], monkeypatch: pytest.MonkeyPatch, ) -> None: admin_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL") if not admin_url: pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL") owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL") runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL") assert owner_url and runtime_url database = "test_onboarding_" + uuid.uuid4().hex restored = "test_onboarding_restore_" + uuid.uuid4().hex admin = create_engine(admin_url, isolation_level="AUTOCOMMIT") owner = make_url(owner_url).set(database=database) runtime = create_engine(make_url(runtime_url).set(database=database)) restored_runtime = create_engine(make_url(runtime_url).set(database=restored)) config = Config() config.set_main_option("script_location", str(Path(__file__).resolve().parents[1] / "alembic")) config.set_main_option("sqlalchemy.url", owner.render_as_string(hide_password=False).replace("%", "%%")) alembic = importlib.import_module("alembic.command") client_id, credential_id = uuid.uuid4(), uuid.uuid4() secret = "mf_" + "p" * 48 with admin.connect() as connection: for name in (database, restored): connection.execute(text(f'CREATE DATABASE "{name}" OWNER modelforge')) try: alembic.upgrade(config, "20260907_0025") with Session(runtime) as session: session.execute(text("INSERT INTO service_clients(id,name,status,allowed_capabilities,requests_per_minute,max_concurrent_requests,workload_priority) VALUES (:id,'prior-client','active','[\"rag.embedding@1\"]',60,1,'interactive')"), {"id": client_id}) session.execute(text("INSERT INTO service_credentials(id,service_client_id,secret_hash,secret_prefix,expires_at) VALUES (:id,:client,:digest,'mf_prior',:expiry)"), {"id": credential_id, "client": client_id, "digest": hashlib.sha256(secret.encode()).hexdigest(), "expiry": datetime.now(UTC) + timedelta(days=30)}) AuditWriter(session, "operator", "prior-release").write("qualification.baseline", "service_client", str(client_id), {}) session.commit() source = semantic_fingerprint(session.connection()) archive = tmp_path / "before-onboarding.dump" args = ["--host", "127.0.0.1" if os.getenv("MODELFORGE_TEST_POSTGRES_CONTAINER") else owner.host or "localhost", "--port", str(owner.port or 5432), "--username", owner.username or "modelforge", "--dbname", database] _database_tool("pg_dump", [*args, "--format=custom"], owner.password, archive) assert archive.stat().st_size > 0 alembic.upgrade(config, "head") with Session(runtime) as session: current = semantic_fingerprint(session.connection()) comparison = fingerprint_compatibility(source, current, source_schema_revision="20260907_0025") assert comparison.compatible assert audit_fingerprint_differences_are_compatible(source, current, comparison) service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None) assert service.authenticate("Bearer " + secret, "rag.embedding@1").id == client_id alembic.downgrade(config, "20260907_0025") alembic.upgrade(config, "head") with Session(runtime) as session: service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None) prior_overlap = service.rotate_client_credential( client_id, ServiceCredentialRotate( expected_credential_id=credential_id, overlap_seconds=1800, expires_at=datetime.now(UTC) + timedelta(days=14), ), ) assert prior_overlap.credential_id prior_overlap_id = prior_overlap.credential_id prior_overlap_secret = prior_overlap.credential prior_overlap_expires_at = prior_overlap.credential_expires_at assert prior_overlap_expires_at assert service.authenticate("Bearer " + secret, "rag.embedding@1").id == client_id with Session(runtime) as session: existing_event_ids = set(session.scalars(select(AuditEvent.id))) # Hold both operations after their current-credential read and before the conditional # claim UPDATE, so the loser must observe the winner's committed CAS change. claim_barrier = threading.Barrier(2) original_claim = ServingService._claim_credential def claim_at_cas(service: ServingService, credential: ServiceCredential) -> None: claim_barrier.wait(timeout=30) original_claim(service, credential) monkeypatch.setattr(ServingService, "_claim_credential", claim_at_cas) race_credential_id = uuid.UUID(str(prior_overlap_id)) expires = datetime.now(UTC) + timedelta(days=14) def contender( action: str, ) -> tuple[str, str, ServiceClientCreated | ServiceClientResponse | None]: with Session(runtime) as session: service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None) try: if action == "rotate": result = service.rotate_client_credential( client_id, ServiceCredentialRotate( expected_credential_id=race_credential_id, overlap_seconds=900, expires_at=expires, ), ) else: result = service.revoke_client_credential( client_id, ServiceCredentialRevoke( expected_credential_id=race_credential_id, confirmation="prior-client", ), ) return action, "applied", result except ServingError as error: assert error.code == "CREDENTIAL_CHANGED" return action, "conflict", None with ThreadPoolExecutor(max_workers=2) as pool: outcomes = list(pool.map(contender, competitors)) assert sorted(outcome for _action, outcome, _result in outcomes) == [ "applied", "conflict", ] winner_action, _, winner_result = next( outcome for outcome in outcomes if outcome[1] == "applied" ) with Session(runtime) as session: new_events = list( session.scalars( select(AuditEvent).where( AuditEvent.action.in_( ["CLIENT_CREDENTIAL_ROTATED", "CLIENT_CREDENTIAL_REVOKED"] ), AuditEvent.id.not_in(existing_event_ids), ) ) ) assert len(new_events) == 1 assert new_events[0].action == ( "CLIENT_CREDENTIAL_ROTATED" if winner_action == "rotate" else "CLIENT_CREDENTIAL_REVOKED" ) credentials = list( session.scalars( select(ServiceCredential).where( ServiceCredential.service_client_id == client_id ) ) ) current = session.get(ServiceCredential, race_credential_id) original = session.get(ServiceCredential, credential_id) assert current and original assert current.superseded_at is not None assert original.superseded_at is not None and original.revoked_at is not None active = [ item for item in credentials if item.superseded_at is None and item.revoked_at is None ] if winner_action == "rotate": assert isinstance(winner_result, ServiceClientCreated) winner_credential_id = uuid.UUID(str(winner_result.credential_id)) assert len(active) == 1 and active[0].id == winner_credential_id assert current.revoked_at is None assert current.expires_at is not None assert datetime.now(UTC) < current.expires_at <= prior_overlap_expires_at assert current.superseded_at is not None assert current.expires_at <= current.superseded_at + timedelta(seconds=900) assert new_events[0].details["previous_credential_id"] == str( race_credential_id ) assert new_events[0].details["credential_id"] == str(winner_credential_id) assert new_events[0].details["overlap_seconds"] == 900 else: assert len(active) == 0 assert current.revoked_at is not None assert all(item.revoked_at is not None for item in credentials) with Session(runtime) as session: service = ServingService(session, Settings(_env_file=None), ManifestRegistry(), None) if winner_action == "rotate": assert isinstance(winner_result, ServiceClientCreated) assert service.authenticate( "Bearer " + prior_overlap_secret, "rag.embedding@1" ).id == client_id assert service.authenticate( "Bearer " + winner_result.credential, "rag.embedding@1" ).id == client_id with pytest.raises(ServingError, match="revoked"): service.authenticate("Bearer " + secret, "rag.embedding@1") else: for retired_secret in (secret, prior_overlap_secret): with pytest.raises(ServingError, match="revoked"): service.authenticate("Bearer " + retired_secret, "rag.embedding@1") with pytest.raises(RuntimeError, match="history exists"): alembic.downgrade(config, "20260907_0025") args[-1] = restored _database_tool("pg_restore", [*args, "--exit-on-error"], owner.password, archive) with restored_runtime.connect() as connection: recovered = semantic_fingerprint(connection) assert recovered["tables"] == source["tables"] config.set_main_option("sqlalchemy.url", owner.set(database=restored).render_as_string(hide_password=False).replace("%", "%%")) alembic.upgrade(config, "head") with restored_runtime.connect() as connection: recovered = semantic_fingerprint(connection) comparison = fingerprint_compatibility(source, recovered, source_schema_revision="20260907_0025") assert comparison.compatible and audit_fingerprint_differences_are_compatible(source, recovered, comparison) finally: runtime.dispose(); restored_runtime.dispose() with admin.connect() as connection: for name in (database, restored): connection.execute(text(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)')) admin.dispose()