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