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

from typing import Any, cast

import pglock
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 structlog.stdlib import get_logger

from django_dramatiq_postgres.conf import Conf
from django_dramatiq_postgres.models import ScheduleBase, TaskState


class Scheduler:
    broker: Broker
    db_alias: str

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.logger = get_logger(__name__, type(self))

    @cached_property
    def model(self) -> type[ScheduleBase]:
        schedule_model = cast(str, Conf().schedule_model)
        model: type[ScheduleBase] = import_string(schedule_model)
        return model

    @property
    def query_set(self) -> QuerySet[ScheduleBase]:
        return self.model._default_manager.using(self.db_alias).filter(paused=False)

    def process_schedule(self, schedule: ScheduleBase) -> None:
        schedule.next_run = schedule.compute_next_run()
        schedule.send(self.broker)
        schedule.save(update_fields=["next_run"])

    def _lock(self) -> pglock.advisory:
        return pglock.advisory(
            lock_id=f"{Conf().channel_prefix}.scheduler",
            side_effect=pglock.Return,
            timeout=0,
        )

    def _run(self) -> int:
        count = 0
        for schedule in self.query_set.filter(next_run__lt=now()):
            if (
                schedule.tasks.using(self.db_alias)  # type: ignore[attr-defined]
                .exclude(state__in=(TaskState.DONE, TaskState.REJECTED))
                .exists()
            ):
                self.logger.debug("Skipping schedule, tasks already exists", schedule=schedule)
                continue
            self.process_schedule(schedule)
            count += 1
        return count

    def run(self) -> int:
        with self._lock() as lock_acquired:
            if not lock_acquired:
                self.logger.debug("Could not acquire lock, skipping scheduling")
                return -1
            count = self._run()
            self.logger.info("Sent scheduled tasks", count=count)
        return count