"""Opt-in real PostgreSQL queue races, using an explicitly disposable database.""" import asyncio import os from pathlib import Path from uuid import uuid4 import pytest from fastapi import HTTPException from psycopg import AsyncConnection from psycopg.conninfo import conninfo_to_dict from psycopg.rows import dict_row from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository def test_concurrent_submission_claims_fencing_and_project_isolation(): dsn = os.getenv("CALCULATION_QUEUE_TEST_DSN") if not dsn: pytest.skip("Set CALCULATION_QUEUE_TEST_DSN to a disposable PostgreSQL database") assert conninfo_to_dict(dsn)["dbname"].startswith("tjwater_calc_test_") async def run(): project, other, owner = uuid4(), uuid4(), uuid4() async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn: migrations = Path(__file__).parents[2] / "resources/sql" await conn.execute((migrations / "026_calculation_jobs.sql").read_text()) legacy_job = uuid4() await conn.execute("""INSERT INTO workbench.calculation_jobs (job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds,status,result,updated_at) VALUES (%s,%s,%s,'legacy','run_simulation','{}','{}',1800,'succeeded','{"preserved":true}',now()-interval '30 days')""", (legacy_job, uuid4(), owner)) lifecycle = (migrations / "027_calculation_job_lifecycle.sql").read_text() await conn.execute(lifecycle) cursor = await conn.execute("SELECT result,result_expires_at,result_expires_at>now()+interval '6 days' AS grace FROM workbench.calculation_jobs WHERE job_id=%s", (legacy_job,)) first = await cursor.fetchone() assert first["result"] == {"preserved": True} and first["grace"] await conn.execute(lifecycle) cursor = await conn.execute("SELECT result_expires_at FROM workbench.calculation_jobs WHERE job_id=%s", (legacy_job,)) assert (await cursor.fetchone())["result_expires_at"] == first["result_expires_at"] async def submit(key, parameters=None, selected_project=project): async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn: return await CalculationJobRepository(conn, selected_project).enqueue( owner_user_id=owner, request_key=key, tool="run_burst_location", parameters=parameters or {"x": 1}, binding={"native_session_id": "session"}, timeout_seconds=600) # Four independent API connections, same logical submission. submissions = await asyncio.gather(*(submit("same") for _ in range(4))) assert len({str(row["job_id"]) for row in submissions}) == 1 with pytest.raises(HTTPException) as conflict: await submit("same", {"x": 2}) assert conflict.value.status_code == 409 second = await submit("second") foreign = await submit("foreign", selected_project=other) async def claim(worker): async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn: return await CalculationJobRepository(conn, project).claim(worker, 60) claimed = await asyncio.gather(*(claim(f"worker-{index}") for index in range(4))) rows = [row for row in claimed if row] assert len(rows) == 2 assert len({str(row["job_id"]) for row in rows}) == 2 assert str(foreign["job_id"]) not in {str(row["job_id"]) for row in rows} async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn: repo = CalculationJobRepository(conn, project) assert await repo.get(foreign["job_id"], owner) is None assert await repo.get(second["job_id"], uuid4()) is None row = rows[0] stale = {**row, "lease_generation": row["lease_generation"] - 1} assert await repo.heartbeat(stale, row["lease_owner"], 60) is None assert not await repo.finish(stale, row["lease_owner"], "succeeded", result={}) cancelled = await repo.cancel(row["job_id"], owner) assert cancelled["status"] == "cancel_requested" assert not await repo.finish(row, row["lease_owner"], "succeeded", result={}) assert await repo.finish(row, row["lease_owner"], "cancelled") expired = rows[1] await conn.execute("UPDATE workbench.calculation_jobs SET lease_expires_at=now()-interval '1 second' WHERE job_id=%s", (expired["job_id"],)) await repo.recover() assert (await repo.get(expired["job_id"], owner))["status"] == "interrupted" assert not await repo.finish(expired, expired["lease_owner"], "succeeded", result={}) assert await repo.claim("replacement", 60) is None # Queued cancellation is terminal immediately, with no later claim. queued = await repo.enqueue(owner_user_id=owner, request_key="cancel-before-start", tool="run_simulation", parameters={}, binding={}, timeout_seconds=30) cancelled = await repo.cancel(queued["job_id"], owner) assert cancelled["status"] == "cancelled" and cancelled["finished_at"] is not None assert cancelled["started_at"] is None assert await repo.claim("replacement", 60) is None # A polling client expires queued jobs without any worker process. waiting = await repo.enqueue(owner_user_id=owner, request_key="expired-queue", tool="run_simulation", parameters={}, binding={}, timeout_seconds=30) await conn.execute("UPDATE workbench.calculation_jobs SET queue_expires_at=now()-interval '1 second' WHERE job_id=%s", (waiting["job_id"],)) status = await repo.get(waiting["job_id"], owner) assert status["status"] == "timed_out" and status["error"]["code"] == "QUEUE_TIMEOUT" assert status["finished_at"] is not None assert "result" not in status # Retention scrubs only payload, keeping deduplication and status. completed = await repo.enqueue(owner_user_id=owner, request_key="retained-identity", tool="run_simulation", parameters={}, binding={}, timeout_seconds=30) await conn.execute("""UPDATE workbench.calculation_jobs SET status='succeeded', result='{"large":"payload"}'::jsonb,finished_at=now(), result_expires_at=now()-interval '1 second' WHERE job_id=%s""", (completed["job_id"],)) expired_result = await repo.get(completed["job_id"], owner, include_result=True) assert expired_result["result"] is None and expired_result["status"] == "succeeded" duplicate = await repo.enqueue(owner_user_id=owner, request_key="retained-identity", tool="run_simulation", parameters={}, binding={}, timeout_seconds=30, max_active_jobs=0) assert duplicate["job_id"] == completed["job_id"] # Concurrent unique requests cannot exceed an otherwise empty project's cap. capacity_project = uuid4() async def bounded_submit(index): try: async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn: return await CalculationJobRepository(conn, capacity_project).enqueue( owner_user_id=owner, request_key=str(index), tool="run_simulation", parameters={}, binding={}, timeout_seconds=30, max_active_jobs=2) except HTTPException as exc: return exc bounded = await asyncio.gather(*(bounded_submit(i) for i in range(8))) assert sum(isinstance(value, dict) for value in bounded) == 2 assert all(isinstance(value, dict) or value.status_code == 503 for value in bounded) asyncio.run(run())