247 lines
12 KiB
Python
247 lines
12 KiB
Python
"""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()
|