facturio/utils/management/commands/runscheduler.py

import logging
import signal
from copy import deepcopy

from apscheduler.schedulers.blocking import BlockingScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from django.conf import settings
from django.core.management.base import BaseCommand
from django_apscheduler import util
from django_apscheduler.jobstores import DjangoJobStore
from django_apscheduler.models import DjangoJobExecution
from django_dramatiq.models import Task

logger = logging.getLogger(__name__)

DEFAULT_JOB_KWARGS = {
    "max_instances": 1,
    "replace_existing": True,
}

DAY_SEC = 24*60*60

PERIODIC_JOBS = [
    {
        "task": "utils.tasks:send_queued_mail_task.send",
        "trigger": IntervalTrigger(seconds=30),
    },

]


# The `close_old_connections` decorator ensures that database connections that have become
# unusable or are obsolete are closed before and after your job has run.
# You should use it to wrap any jobs that you schedule that access the Django database in any way.
@util.close_old_connections
def delete_old_job_executions(max_age: int = 7 * DAY_SEC) -> None:
    """
    This job deletes APScheduler job execution entries older than `max_age` from the database.
    It helps to prevent the database from filling up with old historical records that are no
    longer useful.

    :param max_age: The maximum length of time to retain historical job execution records.
                    Defaults to 7 days.
    """
    DjangoJobExecution.objects.delete_old_job_executions(max_age)
    Task.tasks.delete_old_tasks(max_age)


class Command(BaseCommand):
    help = "Runs APScheduler."  # noqa: A003
    scheduler: BlockingScheduler = None

    def prepare_scheduler(self):
        self.stdout.write(self.style.NOTICE("Preparing scheduler"))
        self.scheduler = BlockingScheduler(timezone=settings.TIME_ZONE)
        self.scheduler.add_jobstore(DjangoJobStore(), "default")

    def add_jobs(self):
        self.scheduler.add_job(
            delete_old_job_executions,
            trigger=CronTrigger(
                day_of_week="mon", hour="00", minute="00",
            ),  # Midnight on Monday, before the start of the next work week.
            id="delete_old_job_executions",
            max_instances=1,
            replace_existing=True,
        )
        self.stdout.write("Added weekly job: 'delete_old_job_executions'.")

        for job in PERIODIC_JOBS:
            kwargs = DEFAULT_JOB_KWARGS | deepcopy(job)
            task = kwargs.pop("task")

            if "id" not in kwargs:
                kwargs["id"] = task

            self.scheduler.add_job(task, **kwargs)

            self.stdout.write(f"Added job: {task}")

    def handle_shutdown(self, *args, **kwargs):  # noqa: ARG001
        self.stdout.write(self.style.NOTICE("Stopping scheduler..."))
        self.scheduler.shutdown()
        self.stdout.write(self.style.NOTICE("Scheduler shut down successfully!"))

    def handle(self, *args, **options):
        self.prepare_scheduler()
        self.add_jobs()

        signal.signal(signal.SIGTERM, self.handle_shutdown)

        try:
            self.stdout.write(self.style.NOTICE("Starting scheduler..."))
            self.scheduler.start()
        except KeyboardInterrupt:
            self.handle_shutdown()