64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
"""Minimal job queue for Agent orchestration: create, atomically claim, complete, fail."""
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import select, update
|
|
|
|
from .models import Job
|
|
|
|
_VALID_KINDS = {"REPROCESS_ALL", "REPROCESS_PRODUCTS"}
|
|
_MAX_ERROR_LENGTH = 2000
|
|
|
|
|
|
def create_job(session, agent, kind, payload=None):
|
|
if kind not in _VALID_KINDS:
|
|
raise ValueError(f"unsupported job kind: {kind}")
|
|
job = Job(agent_id=agent.id, kind=kind, status="PENDING", payload=payload or {})
|
|
session.add(job)
|
|
session.commit()
|
|
return job
|
|
|
|
|
|
def claim_next_job(session, agent):
|
|
"""Atomically claim the oldest PENDING job for this agent only.
|
|
A single UPDATE...WHERE status='PENDING' guards against two concurrent
|
|
pollers (this agent or any other) claiming the same job twice."""
|
|
candidate = session.scalar(
|
|
select(Job.id)
|
|
.where(Job.agent_id == agent.id, Job.status == "PENDING")
|
|
.order_by(Job.created_at.asc())
|
|
.limit(1)
|
|
)
|
|
if candidate is None:
|
|
return None
|
|
now = datetime.now(timezone.utc)
|
|
result = session.execute(
|
|
update(Job)
|
|
.where(Job.id == candidate, Job.status == "PENDING")
|
|
.values(status="RUNNING", started_at=now, attempts=Job.attempts + 1)
|
|
)
|
|
session.commit()
|
|
if result.rowcount != 1:
|
|
return None # another poller won the race; caller can poll again
|
|
return session.get(Job, candidate)
|
|
|
|
|
|
def complete_job(session, job, result_summary):
|
|
result = session.execute(
|
|
update(Job)
|
|
.where(Job.id == job.id, Job.status == "RUNNING")
|
|
.values(status="COMPLETED", finished_at=datetime.now(timezone.utc), result_summary=result_summary)
|
|
)
|
|
session.commit()
|
|
return result.rowcount == 1
|
|
|
|
|
|
def fail_job(session, job, error):
|
|
sanitized = str(error)[:_MAX_ERROR_LENGTH]
|
|
result = session.execute(
|
|
update(Job)
|
|
.where(Job.id == job.id, Job.status == "RUNNING")
|
|
.values(status="FAILED", finished_at=datetime.now(timezone.utc), error=sanitized)
|
|
)
|
|
session.commit()
|
|
return result.rowcount == 1
|