dj-queue 0.1.0__py3-none-any.whl

dj_queue/admin.py

@@ -0,0 +1,90 @@
+import json
+
+from django.contrib import admin, messages
+
+from dj_queue.models import FailedExecution, Job, Pause, Process, RecurringTask, Semaphore
+
+
+@admin.register(Job)
+class JobAdmin(admin.ModelAdmin):
+  list_display = ("id", "task_path", "queue_name", "priority", "status", "created_at")
+  list_select_related = (
+    "ready_execution",
+    "scheduled_execution",
+    "claimed_execution",
+    "blocked_execution",
+    "failed_execution",
+  )
+  readonly_fields = (
+    "task_path",
+    "queue_name",
+    "priority",
+    "payload",
+    "backend_name",
+    "scheduled_at",
+    "concurrency_key",
+    "finished_at",
+    "return_value",
+    "created_at",
+    "updated_at",
+  )
+
+
+@admin.register(FailedExecution)
+class FailedExecutionAdmin(admin.ModelAdmin):
+  list_display = ("job", "exception_class", "message", "created_at")
+  list_select_related = ("job",)
+  actions = ("retry_jobs",)
+  readonly_fields = ("job", "exception_class", "message", "traceback", "created_at")
+
+  @admin.action(description="Retry selected failed jobs")
+  def retry_jobs(self, request, queryset):
+    retried = FailedExecution.retry_all(queryset)
+    self.message_user(request, f"Retried {retried} failed jobs", level=messages.SUCCESS)
+
+
+@admin.register(Process)
+class ProcessAdmin(admin.ModelAdmin):
+  list_display = ("name", "kind", "pid", "hostname", "metadata_json", "last_heartbeat_at")
+  readonly_fields = (
+    "kind",
+    "pid",
+    "hostname",
+    "name",
+    "metadata",
+    "supervisor",
+    "last_heartbeat_at",
+  )
+
+  @admin.display(description="metadata")
+  def metadata_json(self, obj):
+    return json.dumps(obj.metadata, sort_keys=True)
+
+
+@admin.register(RecurringTask)
+class RecurringTaskAdmin(admin.ModelAdmin):
+  list_display = ("key", "task_path", "schedule", "queue_name", "priority", "static")
+  readonly_fields = (
+    "key",
+    "task_path",
+    "payload",
+    "schedule",
+    "queue_name",
+    "priority",
+    "description",
+    "static",
+    "created_at",
+    "updated_at",
+  )
+
+
+@admin.register(Pause)
+class PauseAdmin(admin.ModelAdmin):
+  list_display = ("queue_name", "created_at")
+  readonly_fields = ("queue_name", "created_at")
+
+
+@admin.register(Semaphore)
+class SemaphoreAdmin(admin.ModelAdmin):
+  list_display = ("key", "value", "limit", "expires_at")
+  readonly_fields = ("key", "value", "limit", "expires_at", "created_at", "updated_at")

dj_queue/api.py

@@ -0,0 +1,122 @@
+from functools import partial
+
+from django.db import transaction
+from django.utils import timezone
+from django.utils.module_loading import import_string
+
+from dj_queue.db import get_database_alias
+from dj_queue.log import log_event
+from dj_queue.models import Pause, ReadyExecution, RecurringTask
+
+
+class QueueInfo:
+  def __init__(self, queue_name, *, backend_alias="default"):
+    self.queue_name = queue_name
+    self.backend_alias = backend_alias
+
+  @property
+  def size(self):
+    return self._ready_queryset().count()
+
+  @property
+  def latency(self):
+    oldest = (
+      self._ready_queryset().order_by("created_at").values_list("created_at", flat=True).first()
+    )
+    if oldest is None:
+      return 0.0
+    return (timezone.now() - oldest).total_seconds()
+
+  @property
+  def paused(self):
+    alias = get_database_alias(self.backend_alias)
+    return Pause.objects.using(alias).filter(queue_name=self.queue_name).exists()
+
+  def pause(self):
+    alias = get_database_alias(self.backend_alias)
+    Pause.objects.using(alias).get_or_create(queue_name=self.queue_name)
+    log_event("queue.paused", backend_alias=self.backend_alias, queue_name=self.queue_name)
+
+  def resume(self):
+    alias = get_database_alias(self.backend_alias)
+    deleted, _ = Pause.objects.using(alias).filter(queue_name=self.queue_name).delete()
+    if deleted:
+      log_event("queue.resumed", backend_alias=self.backend_alias, queue_name=self.queue_name)
+
+  def clear(self, *, batch_size=500):
+    deleted = 0
+    while True:
+      job_ids = list(self._ready_queryset().values_list("job_id", flat=True)[:batch_size])
+      if not job_ids:
+        return deleted
+      deleted += _discard_ready_jobs(
+        job_ids=job_ids,
+        batch_size=batch_size,
+        backend_alias=self.backend_alias,
+      )
+
+  @classmethod
+  def all(cls, *, backend_alias="default"):
+    alias = get_database_alias(backend_alias)
+    queue_names = (
+      ReadyExecution.objects.using(alias)
+      .order_by("queue_name")
+      .values_list(
+        "queue_name",
+        flat=True,
+      )
+      .distinct()
+    )
+    return [cls(queue_name, backend_alias=backend_alias) for queue_name in queue_names]
+
+  def _ready_queryset(self):
+    alias = get_database_alias(self.backend_alias)
+    return ReadyExecution.objects.using(alias).filter(queue_name=self.queue_name)
+
+
+def _discard_ready_jobs(*, job_ids, batch_size, backend_alias):
+  discard_ready_jobs = import_string("dj_queue.operations.jobs.discard_ready_jobs")
+  return discard_ready_jobs(job_ids=job_ids, batch_size=batch_size, backend_alias=backend_alias)
+
+
+def enqueue_on_commit(task, *args, using=None, **kwargs):
+  transaction.on_commit(partial(task.enqueue, *args, **kwargs), using=using)
+
+
+def schedule_recurring_task(
+  *,
+  key,
+  task_path,
+  schedule,
+  args=(),
+  kwargs=None,
+  queue_name="default",
+  priority=0,
+  description="",
+  backend_alias="default",
+):
+  alias = get_database_alias(backend_alias)
+  if kwargs is None:
+    kwargs = {}
+
+  recurring_task, _ = RecurringTask.objects.using(alias).update_or_create(
+    key=key,
+    defaults={
+      "task_path": task_path,
+      "payload": {"args": list(args), "kwargs": dict(kwargs)},
+      "schedule": schedule,
+      "queue_name": queue_name,
+      "priority": priority,
+      "description": description,
+      "static": False,
+    },
+  )
+  return recurring_task
+
+
+def unschedule_recurring_task(key, *, backend_alias="default"):
+  alias = get_database_alias(backend_alias)
+  queryset = RecurringTask.objects.using(alias).filter(key=key, static=False)
+  deleted = queryset.count()
+  queryset.delete()
+  return deleted

dj_queue/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class DjQueueConfig(AppConfig):
+  name = "dj_queue"
+  verbose_name = "dj queue"

dj_queue/backend.py

@@ -0,0 +1,161 @@
+from asgiref.sync import sync_to_async
+from django.db import close_old_connections
+from django.tasks import TaskResult, TaskResultStatus
+from django.tasks.backends.base import BaseTaskBackend
+from django.tasks.base import TaskError
+from django.tasks.exceptions import TaskResultDoesNotExist
+from django.utils.module_loading import import_string
+
+from dj_queue.db import get_database_alias
+from dj_queue.models import Job
+from dj_queue.operations.jobs import enqueue_job_with_dispatch, enqueue_jobs_bulk
+
+
+class DjQueueBackend(BaseTaskBackend):
+  supports_async_task = True
+  supports_defer = True
+  supports_get_result = True
+  supports_priority = True
+
+  def enqueue(self, task, args, kwargs):
+    self.validate_task(task)
+    job, dispatched_as = enqueue_job_with_dispatch(task, args, kwargs, backend_alias=self.alias)
+    return _task_result_from_enqueued_job(job, task, dispatched_as)
+
+  async def aenqueue(self, task, args, kwargs):
+    return await sync_to_async(_async_backend_call, thread_sensitive=True)(
+      self.enqueue,
+      task=task,
+      args=args,
+      kwargs=kwargs,
+    )
+
+  def enqueue_all(self, task_calls):
+    jobs = []
+    for task, args, kwargs in task_calls:
+      self.validate_task(task)
+      jobs.append((task, args, kwargs))
+
+    created_jobs = enqueue_jobs_bulk(jobs, backend_alias=self.alias)
+    return [
+      _task_result_from_enqueued_job(job, task, dispatched_as)
+      for job, task, dispatched_as in created_jobs
+    ]
+
+  def get_result(self, result_id):
+    alias = get_database_alias(self.alias)
+    try:
+      job = (
+        Job.objects.using(alias)
+        .select_related(
+          "ready_execution",
+          "scheduled_execution",
+          "claimed_execution__process",
+          "blocked_execution",
+          "failed_execution",
+        )
+        .get(pk=result_id)
+      )
+    except Job.DoesNotExist as exc:
+      raise TaskResultDoesNotExist(str(result_id)) from exc
+
+    return _task_result_from_job(job)
+
+  async def aget_result(self, result_id):
+    return await sync_to_async(_async_backend_call, thread_sensitive=True)(
+      self.get_result,
+      result_id=result_id,
+    )
+
+
+def _async_backend_call(method, /, **kwargs):
+  try:
+    return method(**kwargs)
+  finally:
+    close_old_connections()
+
+
+def _task_result_from_job(job):
+  task = import_string(job.task_path)
+  if hasattr(task, "using"):
+    task = task.using(
+      priority=job.priority,
+      queue_name=job.queue_name,
+      run_after=job.scheduled_at,
+      backend=job.backend_name,
+    )
+
+  status = TaskResultStatus.READY
+  started_at = None
+  finished_at = job.finished_at
+  last_attempted_at = None
+  errors = []
+  worker_ids = []
+
+  if job.failed:
+    status = TaskResultStatus.FAILED
+    finished_at = job.failed_execution.created_at
+    last_attempted_at = job.failed_execution.created_at
+    errors = [
+      TaskError(
+        exception_class_path=job.failed_execution.exception_class,
+        traceback=job.failed_execution.traceback,
+      )
+    ]
+  elif job.claimed:
+    status = TaskResultStatus.RUNNING
+    started_at = job.claimed_execution.created_at
+    last_attempted_at = job.claimed_execution.created_at
+    if job.claimed_execution.process_id is not None:
+      worker_ids = [job.claimed_execution.process.name]
+  elif job.finished:
+    status = TaskResultStatus.SUCCESSFUL
+
+  result = TaskResult(
+    task=task,
+    id=str(job.id),
+    status=status,
+    enqueued_at=job.created_at,
+    started_at=started_at,
+    finished_at=finished_at,
+    last_attempted_at=last_attempted_at,
+    args=job.payload.get("args", []),
+    kwargs=job.payload.get("kwargs", {}),
+    backend=job.backend_name,
+    errors=errors,
+    worker_ids=worker_ids,
+  )
+  if status == TaskResultStatus.SUCCESSFUL:
+    object.__setattr__(result, "_return_value", job.return_value)
+  return result
+
+
+def _task_result_from_enqueued_job(job, task, dispatched_as):
+  if hasattr(task, "using"):
+    task = task.using(
+      priority=job.priority,
+      queue_name=job.queue_name,
+      run_after=job.scheduled_at,
+      backend=job.backend_name,
+    )
+
+  status = TaskResultStatus.SUCCESSFUL if dispatched_as == "discarded" else TaskResultStatus.READY
+  finished_at = job.finished_at if status == TaskResultStatus.SUCCESSFUL else None
+
+  result = TaskResult(
+    task=task,
+    id=str(job.id),
+    status=status,
+    enqueued_at=job.created_at,
+    started_at=None,
+    finished_at=finished_at,
+    last_attempted_at=None,
+    args=job.payload.get("args", []),
+    kwargs=job.payload.get("kwargs", {}),
+    backend=job.backend_name,
+    errors=[],
+    worker_ids=[],
+  )
+  if status == TaskResultStatus.SUCCESSFUL:
+    object.__setattr__(result, "_return_value", job.return_value)
+  return result