feat: add worker job-claim logic
This commit is contained in:
@@ -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
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
import psycopg
|
||||||
|
|
||||||
|
|
||||||
|
def connect() -> psycopg.Connection:
|
||||||
|
dsn = os.environ["DATABASE_URL"]
|
||||||
|
return psycopg.connect(dsn)
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user