fix:registries do not keep connections

Author: cunla

scheduler/helpers/queues/__init__.py

@@ -3,8 +3,8 @@
     "InvalidJobOperation",
     "get_queue",
     "get_all_workers",
-    "perform_job",
+    "queue_perform_job",
 ]
 
 from .getters import get_queue, get_all_workers
-from .queue_logic import Queue, InvalidJobOperation, perform_job
+from .queue_logic import Queue, InvalidJobOperation, queue_perform_job

scheduler/helpers/queues/getters.py

@@ -46,6 +46,12 @@ def _get_connection(config: QueueConfiguration, use_strict_broker: bool = False)
     )
 
 
+def get_queue_connection(queue_name: str) -> ConnectionType:
+    queue_settings = get_queue_configuration(queue_name)
+    connection = _get_connection(queue_settings)
+    return connection
+
+
 def get_queue(name: str = "default") -> Queue:
     """Returns an DjangoQueue using parameters defined in `SCHEDULER_QUEUES`"""
     queue_settings = get_queue_configuration(name)

scheduler/helpers/queues/queue_logic.py

@@ -34,7 +34,7 @@ class NoSuchRegistryError(Exception):
     pass
 
 
-def perform_job(job_model: JobModel, connection: ConnectionType) -> Any:  # noqa
+def queue_perform_job(job_model: JobModel, connection: ConnectionType) -> Any:  # noqa
     """The main execution method. Invokes the job function with the job arguments.
 
     :returns: The job's return value
@@ -271,7 +271,7 @@ def run_sync(self, job: JobModel) -> JobModel:
         """Run a job synchronously, meaning on the same process the method was called."""
         job.prepare_for_execution("sync", self.active_job_registry, self.connection)
         try:
-            result = perform_job(job, self.connection)
+            result = queue_perform_job(job, self.connection)
             self.job_handle_success(job, result=result, job_info_ttl=job.job_info_ttl, result_ttl=job.success_ttl)
         except Exception as e:  # noqa
             logger.warning(f"Job {job.name} failed with exception: {e}")

scheduler/helpers/sentry_integration.py

@@ -32,7 +32,7 @@ def setup_once() -> None:
         version = parse_version(scheduler.__version__)
         _check_minimum_version(SentryIntegration, version)
 
-        old_perform_job = Worker.perform_job
+        old_perform_job = Worker.worker_perform_job
 
         @ensure_integration_enabled(SentryIntegration, old_perform_job)
         def sentry_patched_perform_job(self: Any, job_model: JobModel, *args: Queue, **kwargs: Any) -> bool:
@@ -65,7 +65,7 @@ def sentry_patched_perform_job(self: Any, job_model: JobModel, *args: Queue, **k
 
             return rv
 
-        Worker.perform_job = sentry_patched_perform_job  # type: ignore[method-assign]
+        Worker.worker_perform_job = sentry_patched_perform_job  # type: ignore[method-assign]
 
         old_handle_exception = Worker.handle_exception
 

scheduler/redis_models/registry/base_registry.py

@@ -86,7 +86,7 @@ def key(self) -> str:
 
     @classmethod
     def pop(
-            cls, connection: ConnectionType, registries: Sequence[Self], timeout: Optional[int]
+        cls, connection: ConnectionType, registries: Sequence[Self], timeout: Optional[int]
     ) -> Tuple[Optional[str], Optional[str]]:
         """Helper method to abstract away from some Redis API details
 

scheduler/redis_models/worker.py

@@ -84,12 +84,10 @@ def set_current_job_working_time(self, job_execution_time: float, connection: Co
         self.set_field("current_job_working_time", job_execution_time, connection=connection)
 
     def heartbeat(self, connection: ConnectionType, timeout: Optional[int] = None) -> None:
-        with connection.pipeline() as pipeline:
-            timeout = timeout or DEFAULT_WORKER_TTL + 60
-            pipeline.expire(self._key, timeout)
-            now = utcnow()
-            self.set_field("last_heartbeat", now, connection=pipeline)
-            pipeline.execute()
+        self.last_heartbeat = utcnow()
+        self.save(connection)
+        timeout = timeout or DEFAULT_WORKER_TTL + 60
+        connection.expire(self._key, timeout)
         logger.debug(f"Next heartbeat for worker {self._key} should arrive in {timeout} seconds.")
 
     @classmethod

scheduler/templates/admin/scheduler/jobs.html

@@ -111,7 +111,7 @@
                                 <td>
                                     {{ job.ended_at|date:"Y-m-d, H:i:s" }}
                                 </td>
-                                <td>{{ job.get_status }}</td>
+                                <td>{{ job.status }}</td>
                                 <td>{{ job|show_func_name }}</td>
                                 <td>{{ job.worker_name|default:'-' }}</td>
                                 {% block extra_columns_values %}

scheduler/tests/test_task_types/test_task_model.py

@@ -10,7 +10,7 @@
 
 from scheduler import settings
 from scheduler.helpers.queues import get_queue
-from scheduler.helpers.queues import perform_job
+from scheduler.helpers.queues import queue_perform_job
 from scheduler.models import TaskType, Task, TaskArg, TaskKwarg, run_task
 from scheduler.redis_models import JobStatus, JobModel
 from scheduler.tests import jobs, conf  # noqa
@@ -191,13 +191,13 @@ def test_callable_result(self):
             task = task_factory(self.task_type)
             entry = _get_task_scheduled_job_from_registry(task)
             queue = get_queue("default")
-            self.assertEqual(perform_job(entry, connection=queue.connection), 2)
+            self.assertEqual(queue_perform_job(entry, connection=queue.connection), 2)
 
         def test_callable_empty_args_and_kwargs(self):
             task = task_factory(self.task_type, callable="scheduler.tests.jobs.test_args_kwargs")
             entry = _get_task_scheduled_job_from_registry(task)
             queue = get_queue("default")
-            self.assertEqual(perform_job(entry, connection=queue.connection), "test_args_kwargs()")
+            self.assertEqual(queue_perform_job(entry, connection=queue.connection), "test_args_kwargs()")
 
         def test_delete_args(self):
             task = task_factory(self.task_type)
@@ -244,7 +244,7 @@ def test_callable_args_and_kwargs(self):
             entry = _get_task_scheduled_job_from_registry(task)
             queue = get_queue("default")
             self.assertEqual(
-                perform_job(entry, connection=queue.connection),
+                queue_perform_job(entry, connection=queue.connection),
                 f"test_args_kwargs('one', key1=2, key2={date}, key3=False)",
             )
 

scheduler/tests/test_views/test_queue_actions.py

@@ -106,4 +106,4 @@ def test_job_list_action_stop_jobs__move_to_finished_registry(self):
         self.assertEqual(len(job_names), queue.finished_job_registry.count(queue.connection))
 
         for job_name in job_names:
-            self.assertTrue(queue.finished_job_registry.exists(queue.connection,job_name))
+            self.assertTrue(queue.finished_job_registry.exists(queue.connection, job_name))

scheduler/views/queue_views.py

@@ -84,8 +84,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 = settings.SCHEDULER_CONFIG.TOKEN_VALIDATION_METHOD
-    if request.user.is_staff or (
-            token_validation_func and auth_token and token_validation_func(auth_token)):  # type: ignore
+    if request.user.is_staff or (token_validation_func and auth_token and token_validation_func(auth_token)):  # type: ignore
         return JsonResponse(get_statistics())
 
     return HttpResponseNotFound()