plain.jobs 0.33.0__py3-none-any.whl

plain/jobs/models.py

@@ -0,0 +1,438 @@
+from __future__ import annotations
+
+import datetime
+import logging
+import traceback
+import uuid
+from typing import Self
+
+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
+from plain.runtime import settings
+from plain.utils import timezone
+
+from .registry import jobs_registry
+
+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 = models.DateTimeField(auto_now_add=True)
+    uuid = models.UUIDField(default=uuid.uuid4)
+
+    job_class = models.CharField(max_length=255)
+    parameters = models.JSONField(required=False, allow_null=True)
+    priority = models.IntegerField(default=0)
+    source = models.TextField(required=False)
+    queue = models.CharField(default="default", max_length=255)
+
+    retries = models.IntegerField(default=0)
+    retry_attempt = models.IntegerField(default=0)
+
+    unique_key = models.CharField(max_length=255, required=False)
+
+    start_at = models.DateTimeField(required=False, allow_null=True)
+
+    # OpenTelemetry trace context
+    trace_id = models.CharField(max_length=34, required=False, allow_null=True)
+    span_id = models.CharField(max_length=18, required=False, allow_null=True)
+
+    # expires_at = models.DateTimeField(required=False, allow_null=True)
+
+    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=["unique_key"]),
+            models.Index(fields=["job_class"]),
+            models.Index(fields=["trace_id"]),
+            # Used to dedupe unique in-process jobs
+            models.Index(
+                name="job_request_class_unique_key", fields=["job_class", "unique_key"]
+            ),
+        ],
+        # The job_class and unique_key should be unique at the db-level,
+        # but only if unique_key is not ""
+        constraints=[
+            models.UniqueConstraint(
+                fields=["job_class", "unique_key"],
+                condition=models.Q(unique_key__gt="", retry_attempt=0),
+                name="plainjobs_jobrequest_unique_job_class_key",
+            ),
+            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,
+                unique_key=self.unique_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 = models.UUIDField(default=uuid.uuid4)
+    created_at = models.DateTimeField(auto_now_add=True)
+    started_at = models.DateTimeField(required=False, allow_null=True)
+
+    # From the JobRequest
+    job_request_uuid = models.UUIDField()
+    job_class = models.CharField(max_length=255)
+    parameters = models.JSONField(required=False, allow_null=True)
+    priority = models.IntegerField(default=0)
+    source = models.TextField(required=False)
+    queue = models.CharField(default="default", max_length=255)
+    retries = models.IntegerField(default=0)
+    retry_attempt = models.IntegerField(default=0)
+    unique_key = models.CharField(max_length=255, required=False)
+
+    # OpenTelemetry trace context
+    trace_id = models.CharField(max_length=34, required=False, allow_null=True)
+    span_id = models.CharField(max_length=18, required=False, allow_null=True)
+
+    query = JobQuerySet()
+
+    model_options = models.Options(
+        ordering=["-created_at"],
+        indexes=[
+            models.Index(fields=["created_at"]),
+            models.Index(fields=["queue"]),
+            models.Index(fields=["unique_key"]),
+            models.Index(fields=["started_at"]),
+            models.Index(fields=["job_class"]),
+            models.Index(fields=["job_request_uuid"]),
+            models.Index(fields=["trace_id"]),
+            # Used to dedupe unique in-process jobs
+            models.Index(
+                name="job_class_unique_key", fields=["job_class", "unique_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)
+                job.run()
+                status = JobResultStatuses.SUCCESSFUL
+                error = ""
+                span.set_status(trace.StatusCode.OK)
+            except Exception as e:
+                status = JobResultStatuses.ERRORED
+                error = "".join(traceback.format_tb(e.__traceback__))
+                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=status, error=error)
+
+    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,
+                unique_key=self.unique_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,
+            "unique_key": self.unique_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=models.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"  # Cancelled (probably by deploy)
+    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 = models.UUIDField(default=uuid.uuid4)
+    created_at = models.DateTimeField(auto_now_add=True)
+
+    # From the Job
+    job_process_uuid = models.UUIDField()
+    started_at = models.DateTimeField(required=False, allow_null=True)
+    ended_at = models.DateTimeField(required=False, allow_null=True)
+    error = models.TextField(required=False)
+    status = models.CharField(
+        max_length=20,
+        choices=JobResultStatuses.choices,
+    )
+
+    # From the JobRequest
+    job_request_uuid = models.UUIDField()
+    job_class = models.CharField(max_length=255)
+    parameters = models.JSONField(required=False, allow_null=True)
+    priority = models.IntegerField(default=0)
+    source = models.TextField(required=False)
+    queue = models.CharField(default="default", max_length=255)
+    retries = models.IntegerField(default=0)
+    retry_attempt = models.IntegerField(default=0)
+    unique_key = models.CharField(max_length=255, required=False)
+
+    # Retries
+    retry_job_request_uuid = models.UUIDField(required=False, allow_null=True)
+
+    # OpenTelemetry trace context
+    trace_id = models.CharField(max_length=34, required=False, allow_null=True)
+    span_id = models.CharField(max_length=18, required=False, allow_null=True)
+
+    query = 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"]),
+        ],
+        constraints=[
+            models.UniqueConstraint(
+                fields=["uuid"], name="plainjobs_jobresult_unique_uuid"
+            ),
+        ],
+    )
+
+    def retry_job(self, delay: int | None = None) -> JobRequest:
+        retry_attempt = self.retry_attempt + 1
+        job = jobs_registry.load_job(self.job_class, self.parameters)
+        retry_delay = delay or job.get_retry_delay(retry_attempt)
+
+        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,
+                # Unique key could be passed also?
+            )
+
+            # TODO it is actually possible that result is a list
+            # of pending jobs, which would need to be handled...
+            # Right now it will throw an exception which could be caught by retry_failed_jobs.
+
+            self.retry_job_request_uuid = result.uuid  # type: ignore
+            self.save(update_fields=["retry_job_request_uuid"])
+
+        return result  # type: ignore

plain/jobs/parameters.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import datetime
+from typing import Any
+
+from plain.models import Model, models_registry
+
+
+class JobParameter:
+    """Base class for job parameter serialization/deserialization."""
+
+    STR_PREFIX: str | None = None  # Subclasses should define this
+
+    @classmethod
+    def serialize(cls, value: Any) -> str | None:
+        """Return serialized string or None if can't handle this value."""
+        return None
+
+    @classmethod
+    def deserialize(cls, data: Any) -> Any:
+        """Return deserialized value or None if can't handle this data."""
+        return None
+
+    @classmethod
+    def _extract_string_value(cls, data: Any) -> str | None:
+        """Extract value from string with prefix, return None if invalid format."""
+        if not isinstance(data, str) or not cls.STR_PREFIX:
+            return None
+        if not data.startswith(cls.STR_PREFIX) or len(data) <= len(cls.STR_PREFIX):
+            return None
+        return data[len(cls.STR_PREFIX) :]
+
+
+class ModelParameter(JobParameter):
+    """Handle Plain model instances using a new string format."""
+
+    STR_PREFIX = "__plain://model/"
+
+    @classmethod
+    def serialize(cls, value: Any) -> str | None:
+        if isinstance(value, Model):
+            return f"{cls.STR_PREFIX}{value.model_options.package_label}/{value.model_options.model_name}/{value.id}"
+        return None
+
+    @classmethod
+    def deserialize(cls, data: Any) -> Model | None:
+        if value_part := cls._extract_string_value(data):
+            try:
+                parts = value_part.split("/")
+                if len(parts) == 3 and all(parts):
+                    package, model_name, obj_id = parts
+                    model = models_registry.get_model(package, model_name)
+                    return model.query.get(id=obj_id)
+            except (ValueError, Exception):
+                pass
+        return None
+
+
+class DateParameter(JobParameter):
+    """Handle date objects."""
+
+    STR_PREFIX = "__plain://date/"
+
+    @classmethod
+    def serialize(cls, value: Any) -> str | None:
+        if isinstance(value, datetime.date) and not isinstance(
+            value, datetime.datetime
+        ):
+            return f"{cls.STR_PREFIX}{value.isoformat()}"
+        return None
+
+    @classmethod
+    def deserialize(cls, data: Any) -> datetime.date | None:
+        if value_part := cls._extract_string_value(data):
+            try:
+                return datetime.date.fromisoformat(value_part)
+            except ValueError:
+                pass
+        return None
+
+
+class DateTimeParameter(JobParameter):
+    """Handle datetime objects."""
+
+    STR_PREFIX = "__plain://datetime/"
+
+    @classmethod
+    def serialize(cls, value: Any) -> str | None:
+        if isinstance(value, datetime.datetime):
+            return f"{cls.STR_PREFIX}{value.isoformat()}"
+        return None
+
+    @classmethod
+    def deserialize(cls, data: Any) -> datetime.datetime | None:
+        if value_part := cls._extract_string_value(data):
+            try:
+                return datetime.datetime.fromisoformat(value_part)
+            except ValueError:
+                pass
+        return None
+
+
+class LegacyModelParameter(JobParameter):
+    """Legacy model parameter handling for backwards compatibility."""
+
+    STR_PREFIX = "gid://"
+
+    @classmethod
+    def serialize(cls, value: Any) -> str | None:
+        # Don't serialize new instances with legacy format
+        return None
+
+    @classmethod
+    def deserialize(cls, data: Any) -> Model | None:
+        if value_part := cls._extract_string_value(data):
+            try:
+                package, model, obj_id = value_part.split("/")
+                model = models_registry.get_model(package, model)
+                return model.query.get(id=obj_id)
+            except (ValueError, Exception):
+                pass
+        return None
+
+
+# Registry of parameter types to check in order
+# The order matters - more specific types should come first
+# DateTimeParameter must come before DateParameter since datetime is a subclass of date
+# LegacyModelParameter is last since it only handles deserialization
+PARAMETER_TYPES = [
+    ModelParameter,
+    DateTimeParameter,
+    DateParameter,
+    LegacyModelParameter,
+]
+
+
+class JobParameters:
+    """
+    Main interface for serializing and deserializing job parameters.
+    Uses the registered parameter types to handle different value types.
+    """
+
+    @staticmethod
+    def to_json(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+        serialized_args = []
+        for arg in args:
+            serialized = JobParameters._serialize_value(arg)
+            serialized_args.append(serialized)
+
+        serialized_kwargs = {}
+        for key, value in kwargs.items():
+            serialized = JobParameters._serialize_value(value)
+            serialized_kwargs[key] = serialized
+
+        return {"args": serialized_args, "kwargs": serialized_kwargs}
+
+    @staticmethod
+    def _serialize_value(value: Any) -> Any:
+        """Serialize a single value using the registered parameter types."""
+        # Try each parameter type to see if it can serialize this value
+        for param_type in PARAMETER_TYPES:
+            result = param_type.serialize(value)
+            if result is not None:
+                return result
+
+        # If no parameter type can handle it, return as-is
+        return value
+
+    @staticmethod
+    def from_json(data: dict[str, Any]) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        args = []
+        for arg in data["args"]:
+            deserialized = JobParameters._deserialize_value(arg)
+            args.append(deserialized)
+
+        kwargs = {}
+        for key, value in data["kwargs"].items():
+            deserialized = JobParameters._deserialize_value(value)
+            kwargs[key] = deserialized
+
+        return tuple(args), kwargs
+
+    @staticmethod
+    def _deserialize_value(value: Any) -> Any:
+        """Deserialize a single value using the registered parameter types."""
+        # Try each parameter type to see if it can deserialize this value
+        for param_type in PARAMETER_TYPES:
+            result = param_type.deserialize(value)
+            if result is not None:
+                return result
+
+        # If no parameter type can handle it, return as-is
+        return value

plain/jobs/registry.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from .parameters import JobParameters
+
+if TYPE_CHECKING:
+    from .jobs import Job
+
+T = TypeVar("T", bound=type["Job"])
+
+
+class JobsRegistry:
+    def __init__(self) -> None:
+        self.jobs: dict[str, type[Job]] = {}
+        self.ready = False
+
+    def register_job(self, job_class: type[Job], alias: str = "") -> None:
+        name = self.get_job_class_name(job_class)
+        self.jobs[name] = job_class
+
+        if alias:
+            self.jobs[alias] = job_class
+
+    def get_job_class_name(self, job_class: type[Job]) -> str:
+        return f"{job_class.__module__}.{job_class.__qualname__}"
+
+    def get_job_class(self, name: str) -> type[Job]:
+        return self.jobs[name]
+
+    def load_job(self, job_class_name: str, parameters: dict[str, Any]) -> Job:
+        if not self.ready:
+            raise RuntimeError("Jobs registry is not ready yet")
+
+        job_class = self.get_job_class(job_class_name)
+        args, kwargs = JobParameters.from_json(parameters)
+        return job_class(*args, **kwargs)
+
+
+jobs_registry = JobsRegistry()
+
+
+def register_job(
+    job_class: T | None = None, *, alias: str = ""
+) -> T | Callable[[T], T]:
+    """
+    A decorator that registers a job class in the jobs registry with an optional alias.
+    Can be used both with and without parentheses.
+    """
+    if job_class is None:
+
+        def wrapper(cls: T) -> T:
+            jobs_registry.register_job(cls, alias=alias)
+            return cls
+
+        return wrapper
+    else:
+        jobs_registry.register_job(job_class, alias=alias)
+        return job_class