move schedules to its package

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

packages/django-dramatiq-postgres/django_dramatiq_postgres/models.py

@@ -1,10 +1,18 @@
+import pickle  # nosec
 from enum import StrEnum, auto
 from uuid import uuid4
 
 import pgtrigger
+from cron_converter import Cron
+from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.models import ContentType
+from django.core.exceptions import ValidationError
 from django.db import models
-from django.utils import timezone
+from django.utils.timezone import datetime, now, timedelta
 from django.utils.translation import gettext_lazy as _
+from dramatiq.actor import Actor
+from dramatiq.broker import Broker, get_broker
+from dramatiq.message import Message
 
 from django_dramatiq_postgres.conf import Conf
 
@@ -36,7 +44,7 @@ class TaskBase(models.Model):
         choices=TaskState.choices,
         help_text=_("Task status"),
     )
-    mtime = models.DateTimeField(default=timezone.now, help_text=_("Task last modified time"))
+    mtime = models.DateTimeField(default=now, help_text=_("Task last modified time"))
 
     result = models.BinaryField(null=True, help_text=_("Task result"))
     result_expiry = models.DateTimeField(null=True, help_text=_("Result expiry time"))
@@ -65,3 +73,90 @@ class TaskBase(models.Model):
 
     def __str__(self):
         return str(self.message_id)
+
+
+def validate_crontab(value):
+    try:
+        Cron(value)
+    except ValueError as exc:
+        raise ValidationError(
+            _("%(value)s is not a valid crontab"),
+            params={"value": value},
+        ) from exc
+
+
+class ScheduleBase(models.Model):
+    id = models.UUIDField(primary_key=True, default=uuid4, editable=False)
+
+    actor_name = models.TextField(editable=False, help_text=_("Dramatiq actor to call"))
+    args = models.BinaryField(editable=False, help_text=_("Args to send to the actor"))
+    kwargs = models.BinaryField(editable=False, help_text=_("Kwargs to send to the actor"))
+    options = models.BinaryField(editable=False, help_text=_("Options to send to the actor"))
+
+    rel_obj_content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE, null=True)
+    rel_obj_id = models.TextField(null=True)
+    rel_obj = GenericForeignKey("rel_obj_content_type", "rel_obj_id")
+
+    crontab = models.TextField(validators=[validate_crontab], help_text=_("When to schedule tasks"))
+    paused = models.BooleanField(default=False, help_text=_("Pause this schedule"))
+
+    next_run = models.DateTimeField(auto_now_add=True, editable=False)
+
+    class Meta:
+        abstract = True
+        verbose_name = _("Schedule")
+        verbose_name_plural = _("Schedules")
+        triggers = (
+            pgtrigger.Trigger(
+                name="set_next_run_on_paused",
+                operation=pgtrigger.Update,
+                when=pgtrigger.Before,
+                condition=pgtrigger.Q(new__paused=True) & pgtrigger.Q(old__paused=False),
+                func="""
+                    NEW.next_run = to_timestamp(0);
+                    RETURN NEW;
+                """,
+            ),
+        )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._original_crontab = self.crontab
+
+    def __str__(self):
+        return f"Schedule {self.actor_name} ({self.id})"
+
+    def save(self, *args, **kwargs):
+        if self.crontab != self._original_crontab:
+            self.next_run = self.compute_next_run(now())
+
+        super().save(*args, **kwargs)
+
+        self._original_crontab = self.crontab
+
+    @classmethod
+    def dispatch_by_actor(cls, actor: Actor):
+        """Dispatch a schedule by looking up its actor.
+        Only available for schedules without custom arguments."""
+        schedule = cls.objects.filter(actor_name=actor.actor_name, paused=False).first()
+        if schedule:
+            schedule.send()
+
+    def send(self, broker: Broker | None = None) -> Message:
+        broker = broker or get_broker()
+        actor: Actor = broker.get_actor(self.actor_name)
+        return actor.send_with_options(
+            args=pickle.loads(self.args),  # nosec
+            kwargs=pickle.loads(self.kwargs),  # nosec
+            rel_obj=self,
+            **pickle.loads(self.options),  # nosec
+        )
+
+    def compute_next_run(self, next_run: datetime | None = None) -> datetime:
+        next_run: datetime = self.next_run if not next_run else next_run
+        while True:
+            next_run = Cron(self.crontab).schedule(next_run).next()
+            if next_run > now():
+                return next_run
+            # Force to calculate the one after
+            next_run += timedelta(minutes=1)

packages/django-dramatiq-postgres/django_dramatiq_postgres/scheduler.py

@@ -0,0 +1,50 @@
+import pglock
+from django.db import router, transaction
+from django.db.models import QuerySet
+from django.utils.functional import cached_property
+from django.utils.module_loading import import_string
+from django.utils.timezone import now
+from dramatiq.broker import Broker
+from dramatiq.logging import get_logger
+
+from django_dramatiq_postgres.conf import Conf
+from django_dramatiq_postgres.models import ScheduleBase
+
+
+class Scheduler:
+    def __init__(self, broker: Broker):
+        self.logger = get_logger(__name__, type(self))
+        self.broker = broker
+
+    @cached_property
+    def model(self) -> type[ScheduleBase]:
+        return import_string(Conf().task_class)
+
+    @property
+    def query_set(self) -> QuerySet:
+        return self.model.objects.filter(paused=False)
+
+    def process_schedule(self, schedule: ScheduleBase):
+        schedule.next_run = schedule.compute_next_run()
+        schedule.send(self.broker)
+        schedule.save()
+
+    def _lock(self) -> pglock.advisory:
+        return pglock.advisory(
+            lock_id=f"{Conf().channel_prefix}.scheduler",
+            side_effect=pglock.Return,
+            timeout=0,
+        )
+
+    def _run(self):
+        with transaction.atomic(using=router.db_for_write(self.model)):
+            for schedule in self.query_set.select_for_update().filter(
+                next_run__lt=now(),
+            ):
+                self.process_schedule(schedule)
+
+    def run(self):
+        with self._lock() as lock_acquired:
+            if not lock_acquired:
+                return
+            self._run()