71 خطوط
3.4 KiB
Python
71 خطوط
3.4 KiB
Python
"""Durable PostgreSQL queue (SQLite for local development).
|
|
|
|
Claims use row locks/skip-locked on PostgreSQL and atomic compare-and-swap on
|
|
both engines. The database is the authoritative queue; there is no dual write
|
|
to Redis. Versioned artifact keys fence out expired worker leases.
|
|
"""
|
|
import time
|
|
from functools import lru_cache
|
|
from sqlalchemy import create_engine, String, Text, Float, Integer, select, update, func, text
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
|
|
from .config import settings
|
|
|
|
class Base(DeclarativeBase): pass
|
|
|
|
class Job(Base):
|
|
__tablename__ = 'schematics'
|
|
id: Mapped[str] = mapped_column(String(48), primary_key=True)
|
|
filename: Mapped[str] = mapped_column(String(255))
|
|
status: Mapped[str] = mapped_column(String(20), index=True)
|
|
access: Mapped[str] = mapped_column(String(16))
|
|
token_hash: Mapped[str] = mapped_column(String(64))
|
|
created_at: Mapped[float] = mapped_column(Float)
|
|
updated_at: Mapped[float] = mapped_column(Float)
|
|
expires_at: Mapped[float] = mapped_column(Float, index=True)
|
|
lease_until: Mapped[float | None] = mapped_column(Float, nullable=True)
|
|
attempt: Mapped[int] = mapped_column(Integer, default=0)
|
|
input_key: Mapped[str] = mapped_column(String(200))
|
|
output_prefix: Mapped[str | None] = mapped_column(String(200), nullable=True)
|
|
error_code: Mapped[str | None] = mapped_column(String(60), nullable=True)
|
|
error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
summary: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
|
|
@lru_cache
|
|
def engine():
|
|
url = settings().database_url
|
|
kw = {'connect_args': {'check_same_thread': False, 'timeout': 30}} if url.startswith('sqlite:') else {'pool_pre_ping': True}
|
|
return create_engine(url, **kw)
|
|
|
|
@lru_cache
|
|
def sessions(): return sessionmaker(engine(), expire_on_commit=False)
|
|
|
|
def initialize():
|
|
# Initial schema only. Subsequent schema changes must use explicit migrations.
|
|
if engine().dialect.name == 'postgresql':
|
|
with engine().begin() as conn:
|
|
conn.execute(text('SELECT pg_advisory_xact_lock(84721308)'))
|
|
Base.metadata.create_all(conn)
|
|
else:
|
|
import fcntl
|
|
with open(settings().storage_dir / '.schema.lock', 'a') as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX)
|
|
Base.metadata.create_all(engine())
|
|
|
|
def claim():
|
|
now = time.time()
|
|
with sessions().begin() as db:
|
|
# Expired jobs are never revived. A crashed worker can be retried once.
|
|
eligible = ((Job.status == 'queued') | ((Job.status == 'processing') & (Job.lease_until < now))) & (Job.expires_at > now)
|
|
q = select(Job).where(eligible).order_by(Job.created_at).limit(1)
|
|
if engine().dialect.name == 'postgresql': q = q.with_for_update(skip_locked=True)
|
|
job = db.scalar(q)
|
|
if job is None: return None
|
|
attempt = job.attempt + 1
|
|
values = dict(status='processing', attempt=attempt, updated_at=now,
|
|
lease_until=now + settings().parse_timeout + 60)
|
|
if attempt > 2:
|
|
values.update(status='failed', error_code='worker_lost', error_message='Worker lease expired repeatedly')
|
|
changed = db.execute(update(Job).where(Job.id == job.id, Job.attempt == job.attempt, eligible).values(**values))
|
|
if not changed.rowcount or attempt > 2: return None
|
|
return {'id': job.id, 'filename': job.filename, 'input_key': job.input_key, 'attempt': attempt}
|