plain.jobs 0.43.2__py3-none-any.whl

plain/jobs/models.py

@@ -0,0 +1,567 @@
+from __future__ import annotations
+
+import datetime
+import logging
+import traceback
+from typing import TYPE_CHECKING, Any, Self
+from uuid import UUID, uuid4
+
+from opentelemetry import trace
+from opentelemetry.semconv._incubating.attributes.code_attributes import (
+    CODE_NAMESPACE,
+)
+from opentelemetry.semconv._incubating.attributes.messaging_attributes import (
+    MESSAGING_CONSUMER_GROUP_NAME,
+    MESSAGING_DESTINATION_NAME,
+    MESSAGING_MESSAGE_ID,
+    MESSAGING_OPERATION_NAME,
+    MESSAGING_OPERATION_TYPE,
+    MESSAGING_SYSTEM,
+    MessagingOperationTypeValues,
+)
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
+from opentelemetry.trace import Link, SpanContext, SpanKind
+
+from plain import models
+from plain.models import transaction, types
+from plain.models.expressions import F
+from plain.runtime import settings
+from plain.utils import timezone
+
+from .exceptions import DeferError, DeferJob
+from .registry import jobs_registry
+
+if TYPE_CHECKING:
+    from .jobs import Job
+
+logger = logging.getLogger("plain.jobs")
+tracer = trace.get_tracer("plain.jobs")
+
+
+@models.register_model
+class JobRequest(models.Model):
+    """
+    Keep all pending job requests in a single table.
+    """
+
+    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
+    uuid: UUID = types.UUIDField(default=uuid4)
+
+    job_class: str = types.CharField(max_length=255)
+    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
+    priority: int = types.SmallIntegerField(default=0)
+    source: str = types.TextField(required=False)
+    queue: str = types.CharField(default="default", max_length=255)
+
+    retries: int = types.SmallIntegerField(default=0)
+    retry_attempt: int = types.SmallIntegerField(default=0)
+
+    concurrency_key: str = types.CharField(max_length=255, required=False)
+
+    start_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
+
+    # OpenTelemetry trace context
+    trace_id: str | None = types.CharField(
+        max_length=34, required=False, allow_null=True
+    )
+    span_id: str | None = types.CharField(
+        max_length=18, required=False, allow_null=True
+    )
+
+    # expires_at = models.DateTimeField(required=False, allow_null=True)
+
+    query: models.QuerySet[JobRequest] = models.QuerySet()
+
+    model_options = models.Options(
+        ordering=["priority", "-created_at"],
+        indexes=[
+            models.Index(fields=["priority"]),
+            models.Index(fields=["created_at"]),
+            models.Index(fields=["queue"]),
+            models.Index(fields=["start_at"]),
+            models.Index(fields=["concurrency_key"]),
+            models.Index(fields=["job_class"]),
+            models.Index(fields=["trace_id"]),
+            models.Index(fields=["uuid"]),
+            # Used for job grouping queries
+            models.Index(
+                name="job_request_concurrency_key",
+                fields=["job_class", "concurrency_key"],
+            ),
+        ],
+        constraints=[
+            models.UniqueConstraint(
+                fields=["uuid"], name="plainjobs_jobrequest_unique_uuid"
+            ),
+        ],
+    )
+
+    def __str__(self) -> str:
+        return f"{self.job_class} [{self.uuid}]"
+
+    def convert_to_job_process(self) -> JobProcess:
+        """
+        JobRequests are the pending jobs that are waiting to be executed.
+        We immediately convert them to JobProcess when they are picked up.
+        """
+        with transaction.atomic():
+            result = JobProcess.query.create(
+                job_request_uuid=self.uuid,
+                job_class=self.job_class,
+                parameters=self.parameters,
+                priority=self.priority,
+                source=self.source,
+                queue=self.queue,
+                retries=self.retries,
+                retry_attempt=self.retry_attempt,
+                concurrency_key=self.concurrency_key,
+                trace_id=self.trace_id,
+                span_id=self.span_id,
+            )
+
+            # Delete the pending JobRequest now
+            self.delete()
+
+        return result
+
+
+class JobQuerySet(models.QuerySet["JobProcess"]):
+    def running(self) -> Self:
+        return self.filter(started_at__isnull=False)
+
+    def waiting(self) -> Self:
+        return self.filter(started_at__isnull=True)
+
+    def mark_lost_jobs(self) -> None:
+        # Lost jobs are jobs that have been pending for too long,
+        # and probably never going to get picked up by a worker process.
+        # In theory we could save a timeout per-job and mark them timed-out more quickly,
+        # but if they're still running, we can't actually send a signal to cancel it...
+        now = timezone.now()
+        cutoff = now - datetime.timedelta(seconds=settings.JOBS_TIMEOUT)
+        lost_jobs = self.filter(
+            created_at__lt=cutoff
+        )  # Doesn't matter whether it started or not -- it shouldn't take this long.
+
+        # Note that this will save it in the results,
+        # but lost jobs are only retried if they have a retry!
+        for job in lost_jobs:
+            job.convert_to_result(status=JobResultStatuses.LOST)
+
+
+@models.register_model
+class JobProcess(models.Model):
+    """
+    All active jobs are stored in this table.
+    """
+
+    uuid: UUID = types.UUIDField(default=uuid4)
+    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
+    started_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
+
+    # From the JobRequest
+    job_request_uuid: UUID = types.UUIDField()
+    job_class: str = types.CharField(max_length=255)
+    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
+    priority: int = types.SmallIntegerField(default=0)
+    source: str = types.TextField(required=False)
+    queue: str = types.CharField(default="default", max_length=255)
+    retries: int = types.SmallIntegerField(default=0)
+    retry_attempt: int = types.SmallIntegerField(default=0)
+    concurrency_key: str = types.CharField(max_length=255, required=False)
+
+    # OpenTelemetry trace context
+    trace_id: str | None = types.CharField(
+        max_length=34, required=False, allow_null=True
+    )
+    span_id: str | None = types.CharField(
+        max_length=18, required=False, allow_null=True
+    )
+
+    query: JobQuerySet = JobQuerySet()
+
+    model_options = models.Options(
+        ordering=["-created_at"],
+        indexes=[
+            models.Index(fields=["created_at"]),
+            models.Index(fields=["queue"]),
+            models.Index(fields=["concurrency_key"]),
+            models.Index(fields=["started_at"]),
+            models.Index(fields=["job_class"]),
+            models.Index(fields=["job_request_uuid"]),
+            models.Index(fields=["trace_id"]),
+            models.Index(fields=["uuid"]),
+            # Used for job grouping queries
+            models.Index(
+                name="job_concurrency_key",
+                fields=["job_class", "concurrency_key"],
+            ),
+        ],
+        constraints=[
+            models.UniqueConstraint(fields=["uuid"], name="plainjobs_job_unique_uuid"),
+        ],
+    )
+
+    def run(self) -> JobResult:
+        links = []
+        if self.trace_id and self.span_id:
+            try:
+                links.append(
+                    Link(
+                        SpanContext(
+                            trace_id=int(self.trace_id, 16),
+                            span_id=int(self.span_id, 16),
+                            is_remote=True,
+                        )
+                    )
+                )
+            except (ValueError, TypeError):
+                logger.warning("Invalid trace context for job %s", self.uuid)
+
+        with (
+            tracer.start_as_current_span(
+                f"run {self.job_class}",
+                kind=SpanKind.CONSUMER,
+                attributes={
+                    MESSAGING_SYSTEM: "plain.jobs",
+                    MESSAGING_OPERATION_TYPE: MessagingOperationTypeValues.PROCESS.value,
+                    MESSAGING_OPERATION_NAME: "run",
+                    MESSAGING_MESSAGE_ID: str(self.uuid),
+                    MESSAGING_DESTINATION_NAME: self.queue,
+                    MESSAGING_CONSUMER_GROUP_NAME: self.queue,  # Workers consume from specific queues
+                    CODE_NAMESPACE: self.job_class,
+                },
+                links=links,
+            ) as span
+        ):
+            # This is how we know it has been picked up
+            self.started_at = timezone.now()
+            self.save(update_fields=["started_at"])
+
+            try:
+                job = jobs_registry.load_job(self.job_class, self.parameters or {})
+                job.job_process = self
+
+                try:
+                    job.run()
+                except DeferJob as e:
+                    # Job deferred - not an error, log at INFO level
+                    logger.info(
+                        "Job deferred for %s seconds (increment_retries=%s): job_class=%s job_process_uuid=%s",
+                        e.delay,
+                        e.increment_retries,
+                        self.job_class,
+                        self.uuid,
+                    )
+                    span.set_attribute(ERROR_TYPE, "DeferJob")
+                    span.set_status(trace.StatusCode.OK)  # Not an error
+                    return self.defer(job=job, defer_exception=e)
+
+                # Success case (only reached if no DeferJob was raised)
+                span.set_status(trace.StatusCode.OK)
+                return self.convert_to_result(status=JobResultStatuses.SUCCESSFUL)
+
+            except DeferError as e:
+                # Defer failed (e.g., concurrency limit reached during re-enqueue)
+                # The transaction was rolled back, so the JobProcess still exists in DB.
+                # The pk was restored in defer() before raising, so we can proceed normally.
+                logger.warning("Defer failed for %s: %s", self.job_class, e)
+                span.record_exception(e)
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
+                span.set_attribute(ERROR_TYPE, type(e).__name__)
+                return self.convert_to_result(
+                    status=JobResultStatuses.ERRORED,
+                    error=str(e),
+                )
+
+            except Exception as e:
+                logger.exception(e)
+                span.record_exception(e)
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
+                span.set_attribute(ERROR_TYPE, type(e).__name__)
+                return self.convert_to_result(
+                    status=JobResultStatuses.ERRORED,
+                    error="".join(traceback.format_tb(e.__traceback__)),
+                )
+
+    def defer(self, *, job: Job, defer_exception: DeferJob) -> JobResult:
+        """Defer this job by re-enqueueing it for later execution.
+
+        Atomically deletes the JobProcess, re-enqueues the job, and creates
+        a JobResult linking to the new request. This ensures the concurrency
+        slot is released before attempting to re-enqueue.
+
+        Raises:
+            DeferError: If the job cannot be re-enqueued (e.g., due to concurrency limits).
+                       The transaction will be rolled back and the JobProcess will remain.
+        """
+        # Calculate new retry_attempt based on increment_retries
+        retry_attempt = (
+            self.retry_attempt + 1
+            if defer_exception.increment_retries
+            else self.retry_attempt
+        )
+
+        with transaction.atomic():
+            # 1. Save JobProcess state and delete (releases concurrency slot)
+            saved_id = self.id
+            job_process_uuid = self.uuid
+            job_request_uuid = self.job_request_uuid
+            started_at = self.started_at
+            self.delete()
+
+            # 2. Re-enqueue job (concurrency check can now pass)
+            new_job_request = job.run_in_worker(
+                queue=self.queue,
+                delay=defer_exception.delay,
+                priority=self.priority,
+                retries=self.retries,
+                retry_attempt=retry_attempt,
+                concurrency_key=self.concurrency_key,
+            )
+
+            # Check if re-enqueue failed
+            if new_job_request is None:
+                # Restore id since transaction will roll back and object still exists
+                self.id = saved_id
+                raise DeferError(
+                    f"Failed to re-enqueue deferred job {self.job_class}: "
+                    f"concurrency limit reached for key '{self.concurrency_key}'"
+                )
+
+            # 3. Create JobResult linking to new request
+            result = JobResult.query.create(
+                ended_at=timezone.now(),
+                error=f"Deferred for {defer_exception.delay} seconds",
+                status=JobResultStatuses.DEFERRED,
+                retry_job_request_uuid=new_job_request.uuid,
+                # From the JobProcess
+                job_process_uuid=job_process_uuid,
+                started_at=started_at,
+                # From the JobRequest
+                job_request_uuid=job_request_uuid,
+                job_class=self.job_class,
+                parameters=self.parameters,
+                priority=self.priority,
+                source=self.source,
+                queue=self.queue,
+                retries=self.retries,
+                retry_attempt=self.retry_attempt,
+                concurrency_key=self.concurrency_key,
+                trace_id=self.trace_id,
+                span_id=self.span_id,
+            )
+
+            return result
+
+    def convert_to_result(self, *, status: str, error: str = "") -> JobResult:
+        """
+        Convert this JobProcess to a JobResult.
+        """
+        with transaction.atomic():
+            result = JobResult.query.create(
+                ended_at=timezone.now(),
+                error=error,
+                status=status,
+                # From the JobProcess
+                job_process_uuid=self.uuid,
+                started_at=self.started_at,
+                # From the JobRequest
+                job_request_uuid=self.job_request_uuid,
+                job_class=self.job_class,
+                parameters=self.parameters,
+                priority=self.priority,
+                source=self.source,
+                queue=self.queue,
+                retries=self.retries,
+                retry_attempt=self.retry_attempt,
+                concurrency_key=self.concurrency_key,
+                trace_id=self.trace_id,
+                span_id=self.span_id,
+            )
+
+            # Delete the JobProcess now
+            self.delete()
+
+        return result
+
+    def as_json(self) -> dict[str, str | int | dict | None]:
+        """A JSON-compatible representation to make it easier to reference in Sentry or logging"""
+        return {
+            "uuid": str(self.uuid),
+            "created_at": self.created_at.isoformat(),
+            "started_at": self.started_at.isoformat() if self.started_at else None,
+            "job_request_uuid": str(self.job_request_uuid),
+            "job_class": self.job_class,
+            "parameters": self.parameters,
+            "priority": self.priority,
+            "source": self.source,
+            "queue": self.queue,
+            "retries": self.retries,
+            "retry_attempt": self.retry_attempt,
+            "concurrency_key": self.concurrency_key,
+            "trace_id": self.trace_id,
+            "span_id": self.span_id,
+        }
+
+
+class JobResultQuerySet(models.QuerySet["JobResult"]):
+    def successful(self) -> Self:
+        return self.filter(status=JobResultStatuses.SUCCESSFUL)
+
+    def cancelled(self) -> Self:
+        return self.filter(status=JobResultStatuses.CANCELLED)
+
+    def lost(self) -> Self:
+        return self.filter(status=JobResultStatuses.LOST)
+
+    def errored(self) -> Self:
+        return self.filter(status=JobResultStatuses.ERRORED)
+
+    def retried(self) -> Self:
+        return self.filter(
+            models.Q(retry_job_request_uuid__isnull=False)
+            | models.Q(retry_attempt__gt=0)
+        )
+
+    def failed(self) -> Self:
+        return self.filter(
+            status__in=[
+                JobResultStatuses.ERRORED,
+                JobResultStatuses.LOST,
+                JobResultStatuses.CANCELLED,
+            ]
+        )
+
+    def retryable(self) -> Self:
+        return self.failed().filter(
+            retry_job_request_uuid__isnull=True,
+            retries__gt=0,
+            retry_attempt__lt=F("retries"),
+        )
+
+    def retry_failed_jobs(self) -> None:
+        for result in self.retryable():
+            try:
+                result.retry_job()
+            except Exception:
+                # If something went wrong (like a job class being deleted)
+                # then we immediately increment the retry_attempt on the existing obj
+                # so it won't retry forever.
+                logger.exception(
+                    "Failed to retry job (incrementing retry_attempt): %s", result
+                )
+                result.retry_attempt += 1
+                result.save(update_fields=["retry_attempt"])
+
+
+class JobResultStatuses(models.TextChoices):
+    SUCCESSFUL = "SUCCESSFUL", "Successful"
+    ERRORED = "ERRORED", "Errored"  # Threw an error
+    CANCELLED = "CANCELLED", "Cancelled"  # Interrupted by shutdown/deploy
+    DEFERRED = "DEFERRED", "Deferred"  # Intentionally rescheduled (will run again)
+    LOST = (
+        "LOST",
+        "Lost",
+    )  # Either process lost, lost in transit, or otherwise never finished
+
+
+@models.register_model
+class JobResult(models.Model):
+    """
+    All in-process and completed jobs are stored in this table.
+    """
+
+    uuid: UUID = types.UUIDField(default=uuid4)
+    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
+
+    # From the Job
+    job_process_uuid: UUID = types.UUIDField()
+    started_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
+    ended_at: datetime.datetime | None = types.DateTimeField(
+        required=False, allow_null=True
+    )
+    error: str = types.TextField(required=False)
+    status: str = types.CharField(
+        max_length=20,
+        choices=JobResultStatuses.choices,
+    )
+
+    # From the JobRequest
+    job_request_uuid: UUID = types.UUIDField()
+    job_class: str = types.CharField(max_length=255)
+    parameters: dict[str, Any] | None = types.JSONField(required=False, allow_null=True)
+    priority: int = types.SmallIntegerField(default=0)
+    source: str = types.TextField(required=False)
+    queue: str = types.CharField(default="default", max_length=255)
+    retries: int = types.SmallIntegerField(default=0)
+    retry_attempt: int = types.SmallIntegerField(default=0)
+    concurrency_key: str = types.CharField(max_length=255, required=False)
+
+    # Retries
+    retry_job_request_uuid: UUID | None = types.UUIDField(
+        required=False, allow_null=True
+    )
+
+    # OpenTelemetry trace context
+    trace_id: str | None = types.CharField(
+        max_length=34, required=False, allow_null=True
+    )
+    span_id: str | None = types.CharField(
+        max_length=18, required=False, allow_null=True
+    )
+
+    query: JobResultQuerySet = JobResultQuerySet()
+
+    model_options = models.Options(
+        ordering=["-created_at"],
+        indexes=[
+            models.Index(fields=["created_at"]),
+            models.Index(fields=["job_process_uuid"]),
+            models.Index(fields=["started_at"]),
+            models.Index(fields=["ended_at"]),
+            models.Index(fields=["status"]),
+            models.Index(fields=["job_request_uuid"]),
+            models.Index(fields=["job_class"]),
+            models.Index(fields=["queue"]),
+            models.Index(fields=["trace_id"]),
+            models.Index(fields=["uuid"]),
+        ],
+        constraints=[
+            models.UniqueConstraint(
+                fields=["uuid"], name="plainjobs_jobresult_unique_uuid"
+            ),
+        ],
+    )
+
+    def retry_job(self, delay: int | None = None) -> JobRequest | None:
+        retry_attempt = self.retry_attempt + 1
+        job = jobs_registry.load_job(self.job_class, self.parameters or {})
+
+        if delay is None:
+            retry_delay = job.calculate_retry_delay(retry_attempt)
+        else:
+            retry_delay = delay
+
+        with transaction.atomic():
+            result = job.run_in_worker(
+                # Pass most of what we know through so it stays consistent
+                queue=self.queue,
+                delay=retry_delay,
+                priority=self.priority,
+                retries=self.retries,
+                retry_attempt=retry_attempt,
+                concurrency_key=self.concurrency_key,
+            )
+            if result:
+                self.retry_job_request_uuid = result.uuid
+                self.save(update_fields=["retry_job_request_uuid"])
+                return result
+
+        return None