ruff

Author: cunla

scheduler/helpers/queues/getters.py

@@ -1,14 +1,13 @@
 from typing import List, Set
 
-from scheduler.types import ConnectionErrorTypes, BrokerMetaData
+from scheduler.types import ConnectionErrorTypes, BrokerMetaData, Broker
 from scheduler.redis_models.worker import WorkerModel
 from scheduler.settings import (
     SCHEDULER_CONFIG,
     get_queue_names,
     get_queue_configuration,
     QueueConfiguration,
     logger,
-    Broker,
 )
 from .queue_logic import Queue
 
@@ -22,9 +21,6 @@ class QueueConnectionDiscrepancyError(Exception):
 
 def _get_connection(config: QueueConfiguration, use_strict_broker=False):
     """Returns a Broker connection to use based on parameters in SCHEDULER_QUEUES"""
-    """
-        Returns a redis connection from a connection config
-        """
     if SCHEDULER_CONFIG.BROKER == Broker.FAKEREDIS:
         import fakeredis
 

scheduler/models/task.py

@@ -19,7 +19,8 @@
 from scheduler.helpers.queues import Queue
 from scheduler.helpers.queues import get_queue
 from scheduler.redis_models import JobModel
-from scheduler.settings import logger, get_queue_names, SCHEDULER_CONFIG
+from scheduler.settings import logger, get_queue_names
+from scheduler import settings
 from scheduler.types import ConnectionType, TASK_TYPES
 from .args import TaskArg, TaskKwarg
 from ..helpers import utils
@@ -402,17 +403,18 @@ def clean_queue(self):
             )
 
     def clean_interval_unit(self):
-        if SCHEDULER_CONFIG.SCHEDULER_INTERVAL > self.interval_seconds():
+        config = settings.SCHEDULER_CONFIG
+        if config.SCHEDULER_INTERVAL > self.interval_seconds():
             raise ValidationError(
                 _("Job interval is set lower than %(queue)r queue's interval. minimum interval is %(interval)"),
                 code="invalid",
-                params={"queue": self.queue, "interval": SCHEDULER_CONFIG.SCHEDULER_INTERVAL},
+                params={"queue": self.queue, "interval": config.SCHEDULER_INTERVAL},
             )
-        if self.interval_seconds() % SCHEDULER_CONFIG.SCHEDULER_INTERVAL:
+        if self.interval_seconds() % config.SCHEDULER_INTERVAL:
             raise ValidationError(
                 _("Job interval is not a multiple of rq_scheduler's interval frequency: %(interval)ss"),
                 code="invalid",
-                params={"interval": SCHEDULER_CONFIG.SCHEDULER_INTERVAL},
+                params={"interval": config.SCHEDULER_INTERVAL},
             )
 
     def clean_result_ttl(self) -> None:

scheduler/tests/test_mgmt_commands/test_scheduler_stats.py

@@ -1,10 +1,12 @@
 from django.core.management import call_command
-from django.test import TestCase
+from django.test import TestCase, override_settings
 
 from scheduler.tests import test_settings  # noqa
+from scheduler.types import SchedulerConfiguration
 
 
 class SchedulerStatsTest(TestCase):
+    @override_settings(SCHEDULER_CONFIG=SchedulerConfiguration(SCHEDULER_INTERVAL=1))
     def test_scheduler_stats__does_not_fail(self):
         call_command("scheduler_stats", "-j")
         call_command("scheduler_stats", "-y")

scheduler/tests/test_mgmt_commands/test_scheduler_worker.py

@@ -10,6 +10,7 @@
 
 class SchedulerWorkerTestCase(TestCase):
     def test_scheduler_worker__no_queues_params(self):
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 1
         queue = get_queue("default")
 
         # enqueue some jobs that will fail
@@ -25,6 +26,7 @@ def test_scheduler_worker__no_queues_params(self):
         for job_name in job_names:
             job = JobModel.get(name=job_name, connection=queue.connection)
             self.assertTrue(job.is_failed)
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 10
 
     def test_scheduler_worker__run_jobs(self):
         SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 1
@@ -43,8 +45,10 @@ def test_scheduler_worker__run_jobs(self):
         for job_name in job_names:
             job = JobModel.get(name=job_name, connection=queue.connection)
             self.assertTrue(job.is_failed)
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 10
 
     def test_scheduler_worker__worker_with_two_queues(self):
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 1
         queue = get_queue("default")
         queue2 = get_queue("django_tasks_scheduler_test")
 
@@ -63,8 +67,10 @@ def test_scheduler_worker__worker_with_two_queues(self):
         for job_name in job_names:
             job = JobModel.get(name=job_name, connection=queue.connection)
             self.assertTrue(job.is_failed)
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 10
 
     def test_scheduler_worker__worker_with_one_queue__does_not_perform_other_queue_job(self):
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 1
         queue = get_queue("default")
         queue2 = get_queue("django_tasks_scheduler_test")
 
@@ -80,3 +86,4 @@ def test_scheduler_worker__worker_with_one_queue__does_not_perform_other_queue_j
         other_job = JobModel.get(other_job.name, connection=queue.connection)
 
         self.assertTrue(other_job.is_queued, f"Expected other job to be queued but status={other_job.status}")
+        SCHEDULER_CONFIG.SCHEDULER_INTERVAL = 10

scheduler/tests/test_task_types/test_repeatable_task.py

@@ -9,6 +9,7 @@
 from scheduler.tests.test_task_types.test_task_model import BaseTestCases
 from scheduler.tests.testtools import task_factory, _get_task_scheduled_job_from_registry
 from scheduler.models import TaskType
+from scheduler.types import SchedulerConfiguration
 
 
 class TestRepeatableTask(BaseTestCases.TestSchedulableTask):
@@ -41,7 +42,7 @@ def test_clean_seconds(self):
         job.interval_unit = "seconds"
         self.assertIsNone(job.clean())
 
-    @override_settings(SCHEDULER_CONFIG={"SCHEDULER_INTERVAL": 10})
+    @override_settings(SCHEDULER_CONFIG=SchedulerConfiguration(SCHEDULER_INTERVAL=10))
     def test_clean_too_frequent(self):
         job = task_factory(self.task_type)
         job.queue = self.queue_name
@@ -52,6 +53,7 @@ def test_clean_too_frequent(self):
         with self.assertRaises(ValidationError):
             job.clean_interval_unit()
 
+    @override_settings(SCHEDULER_CONFIG=SchedulerConfiguration(SCHEDULER_INTERVAL=10))
     def test_clean_not_multiple(self):
         job = task_factory(self.task_type)
         job.queue = self.queue_name

scheduler/views/queue_views.py

@@ -14,13 +14,14 @@
 from scheduler.types import ConnectionErrorTypes, ResponseErrorTypes
 from scheduler.helpers.queues import Queue, get_all_workers
 from scheduler.redis_models import JobModel, JobNamesRegistry, WorkerModel
-from scheduler.settings import SCHEDULER_CONFIG, get_queue_names, logger
+from scheduler.settings import get_queue_names, logger
+from scheduler import settings
 from scheduler.views.helpers import get_queue
 from scheduler.worker.commands import StopJobCommand, send_command
 
 
 def _get_registry_job_list(queue: Queue, registry: JobNamesRegistry, page: int) -> Tuple[List[JobModel], int, range]:
-    items_per_page = SCHEDULER_CONFIG.EXECUTIONS_IN_PAGE
+    items_per_page = settings.SCHEDULER_CONFIG.EXECUTIONS_IN_PAGE
     num_jobs = registry.count(queue.connection)
     job_list = list()
 
@@ -85,7 +86,7 @@ def queue_workers(request: HttpRequest, queue_name: str) -> HttpResponse:
 
 def stats_json(request: HttpRequest) -> Union[JsonResponse, HttpResponseNotFound]:
     auth_token = request.headers.get("Authorization")
-    token_validation_func = SCHEDULER_CONFIG.TOKEN_VALIDATION_METHOD
+    token_validation_func = settings.SCHEDULER_CONFIG.TOKEN_VALIDATION_METHOD
     if request.user.is_staff or (token_validation_func and auth_token and token_validation_func(auth_token)):
         return JsonResponse(get_statistics())
 

scheduler/worker/scheduler.py

@@ -24,7 +24,7 @@ class SchedulerStatus(str, Enum):
 
 
 def _reschedule_tasks():
-    enabled_jobs = Task.objects.filter(enabled=True)
+    enabled_jobs = list(Task.objects.filter(enabled=True))
     for item in enabled_jobs:
         logger.debug(f"Rescheduling {str(item)}")
         item.save()
@@ -48,7 +48,7 @@ def __init__(
         self.connection = connection
         self.interval = interval
         self._stop_requested = False
-        self._status = SchedulerStatus.STOPPED
+        self.status = SchedulerStatus.STOPPED
         self._thread = None
         self._pid: Optional[int] = None
         self.worker_name = worker_name
@@ -92,7 +92,7 @@ def start(self) -> None:
         locks = self._acquire_locks()
         if len(locks) == 0:
             return
-        self._status = SchedulerStatus.STARTED
+        self.status = SchedulerStatus.STARTED
         self._thread = Thread(target=run_scheduler, args=(self,), name="scheduler-thread")
         self._thread.start()
 
@@ -117,7 +117,7 @@ def stop(self):
             f"[Scheduler {self.worker_name}/{self.pid}] Stopping scheduler, releasing locks for {', '.join(self._locks.keys())}..."
         )
         self.release_locks()
-        self._status = SchedulerStatus.STOPPED
+        self.status = SchedulerStatus.STOPPED
 
     def release_locks(self):
         """Release acquired locks"""
@@ -145,7 +145,7 @@ def work(self) -> None:
 
     def enqueue_scheduled_jobs(self) -> None:
         """Enqueue jobs whose timestamp is in the past"""
-        self._status = SchedulerStatus.WORKING
+        self.status = SchedulerStatus.WORKING
         _reschedule_tasks()
 
         for registry in self._scheduled_job_registries:
@@ -163,7 +163,7 @@ def enqueue_scheduled_jobs(self) -> None:
                     if job is not None:
                         queue.enqueue_job(job, connection=pipeline, at_front=bool(job.at_front))
                 pipeline.execute()
-        self._status = SchedulerStatus.STARTED
+        self.status = SchedulerStatus.STARTED
 
 
 def run_scheduler(scheduler: WorkerScheduler):

scheduler/worker/worker.py

@@ -31,7 +31,7 @@
     ResponseErrorTypes,
 )
 from .commands import WorkerCommandsChannelListener
-from .scheduler import WorkerScheduler
+from .scheduler import WorkerScheduler, SchedulerStatus
 from ..redis_models.lock import QueueLock
 from ..redis_models.worker import WorkerStatus
 
@@ -337,8 +337,8 @@ def bootstrap(self):
             self.scheduler.start()
             self._model.has_scheduler = True
             self._model.save(connection=self.connection)
-        if self.burst and self.with_scheduler:
-            self.scheduler.stop()
+        if self.with_scheduler and self.burst:
+            self.scheduler.request_stop_and_wait()
             self._model.has_scheduler = False
             self._model.save(connection=self.connection)
         qnames = [queue.name for queue in self.queues]
@@ -371,16 +371,19 @@ def run_maintenance_tasks(self):
         1. Check if scheduler should be started.
         2. Cleaning registries
         """
-        if self.with_scheduler and not self._model.has_scheduler:
+        self.clean_registries()
+        if not self.with_scheduler:
+            return
+        if self.scheduler is None:
             self.scheduler = WorkerScheduler(self.queues, worker_name=self.name, connection=self.connection)
+        if self.scheduler.status == SchedulerStatus.STOPPED:
             self.scheduler.start()
             self._model.has_scheduler = True
             self._model.save(connection=self.connection)
-        if self.burst and self.with_scheduler:
-            self.scheduler.stop()
+        if self.burst:
+            self.scheduler.request_stop_and_wait()
             self._model.has_scheduler = False
             self._model.save(connection=self.connection)
-        self.clean_registries()
 
     def dequeue_job_and_maintain_ttl(
         self, timeout: Optional[int], max_idle_time: Optional[int] = None