156 lines
4.9 KiB
Python
156 lines
4.9 KiB
Python
"""Pry — Background job system using asyncio (built-in) or Celery if available.
|
|
Long-running jobs like crawls, monitors, and bulk imports should not block request threads."""
|
|
|
|
# SPDX-License-Identifier: MIT
|
|
# Copyright (c) 2026 Rug Munch Media LLC
|
|
#
|
|
# Part of Pry — https://git.rugmunch.io/RugMunchMedia/pryscraper
|
|
# Licensed under MIT. See LICENSE.
|
|
import asyncio
|
|
import importlib.util
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from datetime import UTC, datetime
|
|
from enum import StrEnum
|
|
from typing import Any
|
|
|
|
from paths import PRY_DATA_DIR
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Try Celery
|
|
_has_celery = importlib.util.find_spec("celery") is not None
|
|
|
|
JOBS_DIR = PRY_DATA_DIR / "jobs"
|
|
JOBS_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
class JobStatus(StrEnum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
|
|
class Job:
|
|
"""Represents a background job."""
|
|
|
|
def __init__(
|
|
self, job_id: str, name: str, func: Callable, args: tuple = (), kwargs: dict | None = None
|
|
):
|
|
self.id = job_id
|
|
self.name = name
|
|
self.func = func
|
|
self.args = args
|
|
self.kwargs = kwargs or {}
|
|
self.status = JobStatus.PENDING
|
|
self.created_at = datetime.now(UTC).isoformat()
|
|
self.started_at = None
|
|
self.completed_at = None
|
|
self.result = None
|
|
self.error = None
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"id": self.id,
|
|
"name": self.name,
|
|
"status": self.status.value,
|
|
"created_at": self.created_at,
|
|
"started_at": self.started_at,
|
|
"completed_at": self.completed_at,
|
|
"result": self.result,
|
|
"error": self.error,
|
|
}
|
|
|
|
|
|
class JobQueue:
|
|
"""Async background job queue. Uses asyncio for in-process jobs.
|
|
Falls back to Celery if installed and configured."""
|
|
|
|
def __init__(self, max_workers: int = 4, use_celery: bool = False):
|
|
self.max_workers = max_workers
|
|
self.use_celery = use_celery and _has_celery
|
|
self.jobs: dict[str, Job] = {}
|
|
self._queue: asyncio.Queue = asyncio.Queue() if not self.use_celery else None
|
|
self._workers: list[asyncio.Task] = []
|
|
self._running = False
|
|
|
|
async def submit(
|
|
self, name: str, func: Callable, args: tuple = (), kwargs: dict | None = None
|
|
) -> str:
|
|
job_id = uuid.uuid4().hex[:12]
|
|
job = Job(job_id, name, func, args, kwargs)
|
|
self.jobs[job_id] = job
|
|
if self._queue:
|
|
await self._queue.put(job)
|
|
if not self._running:
|
|
await self.start()
|
|
return job_id
|
|
|
|
async def start(self) -> None:
|
|
if self._running or self._queue is None:
|
|
return
|
|
self._running = True
|
|
for i in range(self.max_workers):
|
|
worker = asyncio.create_task(self._worker(f"worker-{i}"))
|
|
self._workers.append(worker)
|
|
|
|
async def _worker(self, name: str) -> None:
|
|
while self._running:
|
|
try:
|
|
job = await asyncio.wait_for(self._queue.get(), timeout=1.0)
|
|
except TimeoutError:
|
|
continue
|
|
await self._execute(job)
|
|
|
|
async def _execute(self, job: Job) -> None:
|
|
job.status = JobStatus.RUNNING
|
|
job.started_at = datetime.now(UTC).isoformat()
|
|
try:
|
|
if asyncio.iscoroutinefunction(job.func):
|
|
job.result = await job.func(*job.args, **job.kwargs)
|
|
else:
|
|
job.result = await asyncio.to_thread(job.func, *job.args, **job.kwargs)
|
|
job.status = JobStatus.COMPLETED
|
|
except Exception as e:
|
|
job.error = str(e)[:500]
|
|
job.status = JobStatus.FAILED
|
|
logger.exception("job_failed", extra={"job_id": job.id, "error": str(e)})
|
|
finally:
|
|
job.completed_at = datetime.now(UTC).isoformat()
|
|
# Persist
|
|
try:
|
|
path = JOBS_DIR / f"{job.id}.json"
|
|
path.write_text(json.dumps(job.to_dict(), indent=2, default=str))
|
|
except OSError:
|
|
pass
|
|
|
|
def get_job(self, job_id: str) -> dict[str, Any] | None:
|
|
job = self.jobs.get(job_id)
|
|
return job.to_dict() if job else None
|
|
|
|
def list_jobs(self, status: str = "", limit: int = 50) -> list[dict[str, Any]]:
|
|
jobs = list(self.jobs.values())
|
|
if status:
|
|
jobs = [j for j in jobs if j.status.value == status]
|
|
return [j.to_dict() for j in sorted(jobs, key=lambda j: j.created_at, reverse=True)[:limit]]
|
|
|
|
async def stop(self) -> None:
|
|
self._running = False
|
|
for w in self._workers:
|
|
w.cancel()
|
|
self._workers.clear()
|
|
|
|
|
|
# Global queue
|
|
_queue: JobQueue | None = None
|
|
|
|
|
|
def get_queue() -> JobQueue:
|
|
global _queue
|
|
if _queue is None:
|
|
_queue = JobQueue()
|
|
return _queue
|