344 lines
11 KiB
Python
344 lines
11 KiB
Python
"""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()]
|