diff --git a/worker/lyra_worker/claim.py b/worker/lyra_worker/claim.py new file mode 100644 index 0000000..0c2f162 --- /dev/null +++ b/worker/lyra_worker/claim.py @@ -0,0 +1,34 @@ +import psycopg + + +def claim_next(conn: psycopg.Connection) -> str | None: + """Atomically claim the oldest queued job. Returns its id or None.""" + with conn.cursor() as cur: + cur.execute( + """ + SELECT id FROM "Job" + WHERE state = 'requested' + ORDER BY "createdAt" ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + """ + ) + row = cur.fetchone() + if row is None: + conn.rollback() + return None + job_id = row[0] + cur.execute( + """ + UPDATE "Job" + SET state = 'matching', + "currentStage" = 'match', + "claimedAt" = now(), + attempts = attempts + 1, + "updatedAt" = now() + WHERE id = %s + """, + (job_id,), + ) + conn.commit() + return job_id diff --git a/worker/lyra_worker/db.py b/worker/lyra_worker/db.py new file mode 100644 index 0000000..8c89282 --- /dev/null +++ b/worker/lyra_worker/db.py @@ -0,0 +1,8 @@ +import os + +import psycopg + + +def connect() -> psycopg.Connection: + dsn = os.environ["DATABASE_URL"] + return psycopg.connect(dsn) diff --git a/worker/tests/conftest.py b/worker/tests/conftest.py new file mode 100644 index 0000000..1b1a662 --- /dev/null +++ b/worker/tests/conftest.py @@ -0,0 +1,36 @@ +import os + +import psycopg +import pytest + + +@pytest.fixture() +def conn(): + dsn = os.environ["DATABASE_URL"] + connection = psycopg.connect(dsn) + yield connection + # Clean up rows created during the test (Job cascades from Request). + with connection.cursor() as cur: + cur.execute('DELETE FROM "Job"') + cur.execute('DELETE FROM "Request"') + connection.commit() + connection.close() + + +def insert_request(conn, artist="Artist", album="Album"): + """Insert a Request + its Job (state 'requested') and return the job id.""" + with conn.cursor() as cur: + cur.execute( + 'INSERT INTO "Request" (id, artist, album, status, "createdAt") ' + "VALUES (gen_random_uuid()::text, %s, %s, 'pending', now()) RETURNING id", + (artist, album), + ) + request_id = cur.fetchone()[0] + cur.execute( + 'INSERT INTO "Job" (id, "requestId", state, "currentStage", attempts, "createdAt", "updatedAt") ' + "VALUES (gen_random_uuid()::text, %s, 'requested', 'intake', 0, now(), now()) RETURNING id", + (request_id,), + ) + job_id = cur.fetchone()[0] + conn.commit() + return job_id diff --git a/worker/tests/test_claim.py b/worker/tests/test_claim.py new file mode 100644 index 0000000..4620247 --- /dev/null +++ b/worker/tests/test_claim.py @@ -0,0 +1,26 @@ +from lyra_worker.claim import claim_next +from tests.conftest import insert_request + + +def test_claim_returns_none_when_empty(conn): + assert claim_next(conn) is None + + +def test_claim_advances_job_to_matching(conn): + job_id = insert_request(conn) + claimed = claim_next(conn) + assert claimed == job_id + + with conn.cursor() as cur: + cur.execute('SELECT state, "currentStage", attempts, "claimedAt" FROM "Job" WHERE id = %s', (job_id,)) + state, stage, attempts, claimed_at = cur.fetchone() + assert state == "matching" + assert stage == "match" + assert attempts == 1 + assert claimed_at is not None + + +def test_claim_does_not_repick_claimed_job(conn): + insert_request(conn) + assert claim_next(conn) is not None + assert claim_next(conn) is None