Build reusable bot framework
This commit is contained in:
Chelsea Lee
343
core/jobs.py
Normal file
343
core/jobs.py
Normal file
@@ -0,0 +1,343 @@
|
||||
"""PostgreSQL-backed scheduled job operations."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
import uuid
|
||||
|
||||
from psycopg2.extras import Json
|
||||
|
||||
from core import postgres
|
||||
|
||||
|
||||
DEFAULT_MAX_ATTEMPTS = 3
|
||||
DEFAULT_LEASE_SECONDS = 300
|
||||
DEFAULT_RETRY_SECONDS = 30
|
||||
MAX_RETRY_SECONDS = 900
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _using_cursor(cursor=None):
|
||||
if cursor is not None:
|
||||
yield cursor
|
||||
return
|
||||
with postgres.get_cursor() as owned_cursor:
|
||||
yield owned_cursor
|
||||
|
||||
|
||||
def _timestamp(value, field="timestamp"):
|
||||
if isinstance(value, str):
|
||||
value = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
if not isinstance(value, datetime):
|
||||
raise ValueError(f"{field} must be a datetime or ISO-8601 string")
|
||||
if value.tzinfo is None:
|
||||
value = value.replace(tzinfo=timezone.utc)
|
||||
return value.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _positive(value, field, maximum=None):
|
||||
try:
|
||||
value = int(value)
|
||||
except (TypeError, ValueError) as error:
|
||||
raise ValueError(f"{field} must be a whole number") from error
|
||||
if value < 1 or (maximum is not None and value > maximum):
|
||||
suffix = f" and at most {maximum}" if maximum is not None else ""
|
||||
raise ValueError(f"{field} must be at least 1{suffix}")
|
||||
return value
|
||||
|
||||
|
||||
def _row(cursor):
|
||||
record = cursor.fetchone()
|
||||
return dict(record) if record else None
|
||||
|
||||
|
||||
def create_job(
|
||||
job_type,
|
||||
payload,
|
||||
run_at,
|
||||
user_uuid=None,
|
||||
max_attempts=DEFAULT_MAX_ATTEMPTS,
|
||||
idempotency_key=None,
|
||||
job_id=None,
|
||||
cursor=None,
|
||||
):
|
||||
"""Create a job, returning the existing row for a repeated idempotency key."""
|
||||
if not isinstance(job_type, str) or not job_type.strip():
|
||||
raise ValueError("job_type is required")
|
||||
max_attempts = _positive(max_attempts, "max_attempts", 100)
|
||||
values = {
|
||||
"id": str(job_id or uuid.uuid4()),
|
||||
"user_uuid": user_uuid,
|
||||
"job_type": job_type.strip(),
|
||||
"payload": Json(payload if payload is not None else {}),
|
||||
"run_at": _timestamp(run_at, "run_at"),
|
||||
"max_attempts": max_attempts,
|
||||
"idempotency_key": idempotency_key,
|
||||
}
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
INSERT INTO scheduled_jobs (
|
||||
id, user_uuid, job_type, payload, run_at,
|
||||
max_attempts, idempotency_key
|
||||
) VALUES (
|
||||
%(id)s, %(user_uuid)s, %(job_type)s, %(payload)s, %(run_at)s,
|
||||
%(max_attempts)s, %(idempotency_key)s
|
||||
)
|
||||
ON CONFLICT (job_type, idempotency_key)
|
||||
DO UPDATE SET idempotency_key = EXCLUDED.idempotency_key
|
||||
WHERE scheduled_jobs.user_uuid IS NOT DISTINCT FROM EXCLUDED.user_uuid
|
||||
RETURNING *
|
||||
""",
|
||||
values,
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def get_job(job_id, cursor=None):
|
||||
"""Return one job by UUID."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute("SELECT * FROM scheduled_jobs WHERE id = %s", (job_id,))
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def list_jobs(user_uuid=None, status=None, job_type=None, limit=100, cursor=None):
|
||||
"""List jobs newest first, optionally filtered by owner, status, or type."""
|
||||
limit = _positive(limit, "limit", 500)
|
||||
clauses = []
|
||||
params = []
|
||||
if user_uuid is not None:
|
||||
clauses.append("user_uuid = %s")
|
||||
params.append(user_uuid)
|
||||
if status is not None:
|
||||
clauses.append("status = %s")
|
||||
params.append(status)
|
||||
if job_type is not None:
|
||||
clauses.append("job_type = %s")
|
||||
params.append(job_type)
|
||||
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
|
||||
params.append(limit)
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
f"""
|
||||
SELECT * FROM scheduled_jobs
|
||||
{where}
|
||||
ORDER BY created_at DESC
|
||||
LIMIT %s
|
||||
""",
|
||||
params,
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
|
||||
|
||||
def claim_due_jobs(
|
||||
worker_id,
|
||||
limit=10,
|
||||
lease_seconds=DEFAULT_LEASE_SECONDS,
|
||||
job_types=None,
|
||||
cursor=None,
|
||||
):
|
||||
"""Atomically lease due jobs using row locks that skip other workers."""
|
||||
if not worker_id:
|
||||
raise ValueError("worker_id is required")
|
||||
limit = _positive(limit, "limit", 100)
|
||||
lease_seconds = _positive(lease_seconds, "lease_seconds")
|
||||
if isinstance(job_types, str):
|
||||
job_types = [job_types]
|
||||
elif job_types is not None:
|
||||
job_types = list(job_types)
|
||||
if not job_types:
|
||||
return []
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
WITH exhausted AS (
|
||||
SELECT id
|
||||
FROM scheduled_jobs
|
||||
WHERE status = 'running'
|
||||
AND (lease_until IS NULL OR lease_until <= NOW())
|
||||
AND attempts >= max_attempts
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE scheduled_jobs AS job
|
||||
SET status = 'failed',
|
||||
lease_until = NULL,
|
||||
leased_by = NULL,
|
||||
last_error = COALESCE(last_error, 'lease expired'),
|
||||
updated_at = NOW()
|
||||
FROM exhausted
|
||||
WHERE job.id = exhausted.id
|
||||
"""
|
||||
)
|
||||
active_cursor.execute(
|
||||
"""
|
||||
WITH candidates AS (
|
||||
SELECT id
|
||||
FROM scheduled_jobs
|
||||
WHERE run_at <= NOW()
|
||||
AND attempts < max_attempts
|
||||
AND (
|
||||
status = 'pending'
|
||||
OR (
|
||||
status = 'running'
|
||||
AND (lease_until IS NULL OR lease_until <= NOW())
|
||||
)
|
||||
)
|
||||
AND (%(job_types)s IS NULL OR job_type = ANY(%(job_types)s))
|
||||
ORDER BY run_at, created_at
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT %(limit)s
|
||||
)
|
||||
UPDATE scheduled_jobs AS job
|
||||
SET status = 'running',
|
||||
attempts = job.attempts + 1,
|
||||
leased_by = %(worker_id)s,
|
||||
lease_until = NOW() + (%(lease_seconds)s * INTERVAL '1 second'),
|
||||
updated_at = NOW()
|
||||
FROM candidates
|
||||
WHERE job.id = candidates.id
|
||||
RETURNING job.*
|
||||
""",
|
||||
{
|
||||
"job_types": job_types,
|
||||
"limit": limit,
|
||||
"worker_id": worker_id,
|
||||
"lease_seconds": lease_seconds,
|
||||
},
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
|
||||
|
||||
def renew_job_lease(job_id, worker_id, lease_seconds=DEFAULT_LEASE_SECONDS, cursor=None):
|
||||
"""Extend a lease only while it is owned by the requesting worker."""
|
||||
lease_seconds = _positive(lease_seconds, "lease_seconds")
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET lease_until = NOW() + (%s * INTERVAL '1 second'), updated_at = NOW()
|
||||
WHERE id = %s AND status = 'running' AND leased_by = %s
|
||||
RETURNING *
|
||||
""",
|
||||
(lease_seconds, job_id, worker_id),
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def complete_job(job_id, worker_id, cursor=None):
|
||||
"""Mark a job complete when its lease is still owned by the worker."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'completed', completed_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL, last_error = NULL
|
||||
WHERE id = %s AND status = 'running' AND leased_by = %s
|
||||
RETURNING *
|
||||
""",
|
||||
(job_id, worker_id),
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def fail_job(
|
||||
job_id,
|
||||
worker_id,
|
||||
error,
|
||||
retry_seconds=DEFAULT_RETRY_SECONDS,
|
||||
max_retry_seconds=MAX_RETRY_SECONDS,
|
||||
cursor=None,
|
||||
):
|
||||
"""Fail or reschedule an owned job using bounded exponential backoff."""
|
||||
retry_seconds = _positive(retry_seconds, "retry_seconds")
|
||||
max_retry_seconds = _positive(max_retry_seconds, "max_retry_seconds")
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
SELECT attempts, max_attempts
|
||||
FROM scheduled_jobs
|
||||
WHERE id = %s AND status = 'running' AND leased_by = %s
|
||||
FOR UPDATE
|
||||
""",
|
||||
(job_id, worker_id),
|
||||
)
|
||||
current = active_cursor.fetchone()
|
||||
if not current:
|
||||
return None
|
||||
|
||||
exhausted = current["attempts"] >= current["max_attempts"]
|
||||
delay = min(
|
||||
max_retry_seconds,
|
||||
retry_seconds * (2 ** min(max(current["attempts"] - 1, 0), 30)),
|
||||
)
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = %(status)s,
|
||||
run_at = CASE
|
||||
WHEN %(exhausted)s THEN run_at
|
||||
ELSE NOW() + (%(delay)s * INTERVAL '1 second')
|
||||
END,
|
||||
leased_by = NULL,
|
||||
lease_until = NULL,
|
||||
last_error = %(error)s,
|
||||
updated_at = NOW()
|
||||
WHERE id = %(job_id)s
|
||||
RETURNING *
|
||||
""",
|
||||
{
|
||||
"status": "failed" if exhausted else "pending",
|
||||
"exhausted": exhausted,
|
||||
"delay": delay,
|
||||
"error": str(error)[:4000],
|
||||
"job_id": job_id,
|
||||
},
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def cancel_job(job_id, user_uuid=None, cursor=None):
|
||||
"""Cancel one unfinished job, optionally enforcing its owner."""
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE id = %(job_id)s
|
||||
AND status IN ('pending', 'running')
|
||||
AND (%(user_uuid)s IS NULL OR user_uuid = %(user_uuid)s)
|
||||
RETURNING *
|
||||
""",
|
||||
{"job_id": job_id, "user_uuid": user_uuid},
|
||||
)
|
||||
return _row(active_cursor)
|
||||
|
||||
|
||||
def cancel_jobs(
|
||||
user_uuid=None, job_type=None, idempotency_key=None, cursor=None
|
||||
):
|
||||
"""Cancel matching unfinished jobs; at least one filter is required."""
|
||||
filters = {
|
||||
"user_uuid": user_uuid,
|
||||
"job_type": job_type,
|
||||
"idempotency_key": idempotency_key,
|
||||
}
|
||||
clauses = [f"{name} = %({name})s" for name, value in filters.items() if value is not None]
|
||||
if not clauses:
|
||||
raise ValueError("at least one cancellation filter is required")
|
||||
|
||||
with _using_cursor(cursor) as active_cursor:
|
||||
active_cursor.execute(
|
||||
f"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE status IN ('pending', 'running')
|
||||
AND {' AND '.join(clauses)}
|
||||
RETURNING *
|
||||
""",
|
||||
filters,
|
||||
)
|
||||
return [dict(record) for record in active_cursor.fetchall()]
|
||||
Reference in New Issue
Block a user