230 lines
7.2 KiB
Python
230 lines
7.2 KiB
Python
|
from collections.abc import Iterable
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from types import TracebackType
|
||
|
|
from typing import Any, Optional, TypeVar
|
||
|
|
|
||
|
|
import django_rq
|
||
|
|
from django.apps import apps
|
||
|
|
from django.core.checks import messages
|
||
|
|
from django.core.exceptions import SuspiciousOperation
|
||
|
|
from django.db import transaction
|
||
|
|
from django.utils.module_loading import import_string
|
||
|
|
from redis.client import Redis
|
||
|
|
from rq.job import Callback, JobStatus
|
||
|
|
from rq.job import Job as BaseJob
|
||
|
|
from rq.registry import ScheduledJobRegistry
|
||
|
|
from typing_extensions import ParamSpec
|
||
|
|
|
||
|
|
from django_tasks.backends.base import BaseTaskBackend
|
||
|
|
from django_tasks.exceptions import ResultDoesNotExist
|
||
|
|
from django_tasks.signals import task_enqueued, task_finished
|
||
|
|
from django_tasks.task import DEFAULT_PRIORITY, MAX_PRIORITY, ResultStatus, Task
|
||
|
|
from django_tasks.task import TaskResult as BaseTaskResult
|
||
|
|
from django_tasks.utils import get_module_path, get_random_id
|
||
|
|
|
||
|
|
T = TypeVar("T")
|
||
|
|
P = ParamSpec("P")
|
||
|
|
|
||
|
|
RQ_STATUS_TO_RESULT_STATUS = {
|
||
|
|
JobStatus.QUEUED: ResultStatus.NEW,
|
||
|
|
JobStatus.FINISHED: ResultStatus.SUCCEEDED,
|
||
|
|
JobStatus.FAILED: ResultStatus.FAILED,
|
||
|
|
JobStatus.STARTED: ResultStatus.RUNNING,
|
||
|
|
JobStatus.DEFERRED: ResultStatus.NEW,
|
||
|
|
JobStatus.SCHEDULED: ResultStatus.NEW,
|
||
|
|
JobStatus.STOPPED: ResultStatus.FAILED,
|
||
|
|
JobStatus.CANCELED: ResultStatus.FAILED,
|
||
|
|
None: ResultStatus.NEW,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class TaskResult(BaseTaskResult[T]):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class Job(BaseJob):
|
||
|
|
def _execute(self) -> Any:
|
||
|
|
"""
|
||
|
|
Shim RQ's `Job` to call the underlying `Task` function.
|
||
|
|
"""
|
||
|
|
return self.func.call(*self.args, **self.kwargs)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def func(self) -> Task:
|
||
|
|
func = super().func
|
||
|
|
|
||
|
|
if not isinstance(func, Task):
|
||
|
|
raise SuspiciousOperation(
|
||
|
|
f"Task {self.id} does not point to a Task ({self.func_name})"
|
||
|
|
)
|
||
|
|
|
||
|
|
return func
|
||
|
|
|
||
|
|
def into_task_result(self) -> TaskResult:
|
||
|
|
task: Task = self.func
|
||
|
|
|
||
|
|
scheduled_job_registry = ScheduledJobRegistry( # type: ignore[no-untyped-call]
|
||
|
|
queue=django_rq.get_queue(self.origin)
|
||
|
|
)
|
||
|
|
|
||
|
|
if self.is_scheduled:
|
||
|
|
run_after = scheduled_job_registry.get_scheduled_time(self)
|
||
|
|
else:
|
||
|
|
run_after = None
|
||
|
|
|
||
|
|
task_result: TaskResult = TaskResult(
|
||
|
|
task=task.using(
|
||
|
|
priority=DEFAULT_PRIORITY,
|
||
|
|
queue_name=self.origin,
|
||
|
|
run_after=run_after,
|
||
|
|
backend=self.meta["backend_name"],
|
||
|
|
),
|
||
|
|
id=self.id,
|
||
|
|
status=RQ_STATUS_TO_RESULT_STATUS[self.get_status()],
|
||
|
|
enqueued_at=self.enqueued_at,
|
||
|
|
started_at=self.started_at,
|
||
|
|
finished_at=self.ended_at,
|
||
|
|
args=list(self.args),
|
||
|
|
kwargs=self.kwargs,
|
||
|
|
backend=self.meta["backend_name"],
|
||
|
|
)
|
||
|
|
|
||
|
|
latest_result = self.latest_result()
|
||
|
|
|
||
|
|
if latest_result is not None:
|
||
|
|
if "exception_class" in self.meta:
|
||
|
|
object.__setattr__(
|
||
|
|
task_result,
|
||
|
|
"_exception_class",
|
||
|
|
import_string(self.meta["exception_class"]),
|
||
|
|
)
|
||
|
|
object.__setattr__(task_result, "_traceback", latest_result.exc_string)
|
||
|
|
object.__setattr__(task_result, "_return_value", latest_result.return_value)
|
||
|
|
|
||
|
|
return task_result
|
||
|
|
|
||
|
|
|
||
|
|
def failed_callback(
|
||
|
|
job: Job,
|
||
|
|
connection: Optional[Redis],
|
||
|
|
exception_class: type[Exception],
|
||
|
|
exception_value: Exception,
|
||
|
|
traceback: TracebackType,
|
||
|
|
) -> None:
|
||
|
|
task_result = job.into_task_result()
|
||
|
|
|
||
|
|
# Smuggle the exception class through meta
|
||
|
|
job.meta["exception_class"] = get_module_path(exception_class)
|
||
|
|
job.save_meta() # type: ignore[no-untyped-call]
|
||
|
|
|
||
|
|
task_finished.send(type(task_result.task.get_backend()), task_result=task_result)
|
||
|
|
|
||
|
|
|
||
|
|
def success_callback(job: Job, connection: Optional[Redis], result: Any) -> None:
|
||
|
|
task_result = job.into_task_result()
|
||
|
|
|
||
|
|
task_finished.send(type(task_result.task.get_backend()), task_result=task_result)
|
||
|
|
|
||
|
|
|
||
|
|
class RQBackend(BaseTaskBackend):
|
||
|
|
supports_async_task = True
|
||
|
|
supports_get_result = True
|
||
|
|
supports_defer = True
|
||
|
|
|
||
|
|
def __init__(self, alias: str, params: dict) -> None:
|
||
|
|
super().__init__(alias, params)
|
||
|
|
|
||
|
|
if not self.queues:
|
||
|
|
self.queues = set(django_rq.settings.QUEUES.keys())
|
||
|
|
|
||
|
|
def enqueue(
|
||
|
|
self,
|
||
|
|
task: Task[P, T],
|
||
|
|
args: P.args, # type:ignore[valid-type]
|
||
|
|
kwargs: P.kwargs, # type:ignore[valid-type]
|
||
|
|
) -> TaskResult[T]:
|
||
|
|
self.validate_task(task)
|
||
|
|
|
||
|
|
queue = django_rq.get_queue(task.queue_name, job_class=Job)
|
||
|
|
|
||
|
|
task_result = TaskResult[T](
|
||
|
|
task=task,
|
||
|
|
id=get_random_id(),
|
||
|
|
status=ResultStatus.NEW,
|
||
|
|
enqueued_at=None,
|
||
|
|
started_at=None,
|
||
|
|
finished_at=None,
|
||
|
|
args=args,
|
||
|
|
kwargs=kwargs,
|
||
|
|
backend=self.alias,
|
||
|
|
)
|
||
|
|
|
||
|
|
job = queue.create_job(
|
||
|
|
task.module_path,
|
||
|
|
args=args,
|
||
|
|
kwargs=kwargs,
|
||
|
|
job_id=task_result.id,
|
||
|
|
status=JobStatus.SCHEDULED if task.run_after else JobStatus.QUEUED,
|
||
|
|
on_failure=Callback(failed_callback),
|
||
|
|
on_success=Callback(success_callback),
|
||
|
|
meta={"backend_name": self.alias},
|
||
|
|
)
|
||
|
|
|
||
|
|
def save_result() -> None:
|
||
|
|
nonlocal job
|
||
|
|
if task.run_after is None:
|
||
|
|
job = queue.enqueue_job(job, at_front=task.priority == MAX_PRIORITY)
|
||
|
|
else:
|
||
|
|
job = queue.schedule_job(job, task.run_after)
|
||
|
|
|
||
|
|
object.__setattr__(task_result, "enqueued_at", job.enqueued_at)
|
||
|
|
|
||
|
|
task_enqueued.send(type(self), task_result=task_result)
|
||
|
|
|
||
|
|
if self._get_enqueue_on_commit_for_task(task):
|
||
|
|
transaction.on_commit(save_result)
|
||
|
|
else:
|
||
|
|
save_result()
|
||
|
|
|
||
|
|
return task_result
|
||
|
|
|
||
|
|
def _get_queues(self) -> list[django_rq.queues.DjangoRQ]:
|
||
|
|
return django_rq.queues.get_queues(*self.queues, job_class=Job) # type: ignore[no-any-return]
|
||
|
|
|
||
|
|
def _get_job(self, job_id: str) -> Optional[Job]:
|
||
|
|
for queue in self._get_queues():
|
||
|
|
job = queue.fetch_job(job_id)
|
||
|
|
if job is not None:
|
||
|
|
return job # type: ignore[no-any-return]
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_result(self, result_id: str) -> TaskResult:
|
||
|
|
job = self._get_job(result_id)
|
||
|
|
|
||
|
|
if job is None:
|
||
|
|
raise ResultDoesNotExist(result_id)
|
||
|
|
|
||
|
|
return job.into_task_result()
|
||
|
|
|
||
|
|
def check(self, **kwargs: Any) -> Iterable[messages.CheckMessage]:
|
||
|
|
yield from super().check(**kwargs)
|
||
|
|
|
||
|
|
backend_name = self.__class__.__name__
|
||
|
|
|
||
|
|
if not apps.is_installed("django_rq"):
|
||
|
|
yield messages.Error(
|
||
|
|
f"{backend_name} configured as django_tasks backend, but django_rq app not installed",
|
||
|
|
"Insert 'django_rq' in INSTALLED_APPS",
|
||
|
|
)
|
||
|
|
|
||
|
|
for queue_name in self.queues:
|
||
|
|
try:
|
||
|
|
django_rq.get_queue(queue_name)
|
||
|
|
except KeyError:
|
||
|
|
yield messages.Error(
|
||
|
|
f"{queue_name!r} is not configured for django-rq",
|
||
|
|
f"Add {queue_name!r} to RQ_QUEUES",
|
||
|
|
)
|