Wrap tasks in app executor

Author: knifecake

Makefile

@@ -1,9 +1,13 @@
+.PHONY: test
 test:
+	tests/bin/run-docker-postgres
 	uv run python runtests.py
 
-test.postgres:
+.PHONY: steady_queue
+steady_queue:
 	tests/bin/run-docker-postgres
-	DB_URL=postgres://steady_queue:steady_queue@localhost:5432/steady_queue uv run python runtests.py
+	uv run python manage.py steady_queue
 
+.PHONY: lint
 lint:
 	uv run pre-commit run --all-files

steady_queue/__init__.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 from typing import Optional
 
-VERSION = (0, 1, "0-beta2")
+VERSION = (0, 1, "0b3")
 
 __version__ = ".".join(map(str, VERSION))
 

steady_queue/app_executor.py

@@ -0,0 +1,45 @@
+import logging
+import threading
+from contextlib import contextmanager
+
+from django.db import connections
+
+logger = logging.getLogger("steady_queue")
+
+
+class AppExecutor:
+    @staticmethod
+    @contextmanager
+    def wrap_in_app_executor():
+        """
+        Django equivalent of Rails' wrap_in_app_executor.
+        Ensures proper database connection handling in background threads.
+        """
+        # Ensure we have a database connection for this thread
+        connection = connections["default"]  # or your specific database alias
+
+        try:
+            # Ensure connection is established
+            connection.ensure_connection()
+
+            # Clear any existing queries from connection (equivalent to Rails query cache reset)
+            # if hasattr(connection, "queries_logged"):
+            #     connection.queries_logged = 0
+            #     connection.queries = []
+
+            yield
+
+        except Exception as e:
+            # Handle any database-related errors
+            # Close connection on error to prevent connection leaks
+            logger.exception("Error in AppExecutor: %s", e)
+            if connection.connection:
+                connection.close()
+            raise
+        finally:
+            # Django automatically manages connection lifecycle,
+            # but we can explicitly close if needed for long-running processes
+            if threading.current_thread() != threading.main_thread():
+                # Only close connections in background threads
+                # Main thread connections are handled by Django's request/response cycle
+                connection.close()

steady_queue/processes/base.py

@@ -3,6 +3,8 @@
 import socket
 from typing import Any
 
+from django.db import connections
+
 
 class Base:
     name: str
@@ -43,3 +45,6 @@ def is_stopped(self) -> bool:
 
     def generate_name(self) -> str:
         return "-".join((self.kind, secrets.token_hex(10)))
+
+    def reset_database_connections(self):
+        connections.close_all()

steady_queue/processes/maintenance.py

@@ -1,6 +1,7 @@
 import logging
 
 import steady_queue
+from steady_queue.app_executor import AppExecutor
 from steady_queue.models.claimed_execution import ClaimedExecution
 from steady_queue.models.process import Process
 from steady_queue.processes.errors import ProcessMissingError
@@ -23,8 +24,10 @@ def stop_maintenance_task(self):
         self.maintenance_task.stop()
 
     def fail_orphaned_executions(self):
-        ClaimedExecution.objects.orphaned().fail_all_with(ProcessMissingError())
+        with AppExecutor.wrap_in_app_executor():
+            ClaimedExecution.objects.orphaned().fail_all_with(ProcessMissingError())
 
     def prune_dead_processes(self):
         logger.debug("pruning dead processes")
-        Process.objects.exclude(pk=self.process.pk).prune()
+        with AppExecutor.wrap_in_app_executor():
+            Process.objects.exclude(pk=self.process.pk).prune()

steady_queue/processes/poller.py

@@ -2,6 +2,7 @@
 from datetime import timedelta
 from typing import Any
 
+from steady_queue.app_executor import AppExecutor
 from steady_queue.processes.base import Base
 from steady_queue.processes.interruptible import Interruptible
 from steady_queue.processes.registrable import Registrable
@@ -38,7 +39,9 @@ def start_loop(self):
                 if self.is_shutting_down:
                     break
 
-                delay = self.poll()
+                with AppExecutor.wrap_in_app_executor():
+                    delay = self.poll()
+
                 self.interruptible_sleep(delay)
         finally:
             self.shutdown()

steady_queue/processes/pool.py

@@ -3,6 +3,7 @@
 from threading import Lock
 from typing import Callable
 
+from steady_queue.app_executor import AppExecutor
 from steady_queue.models.claimed_execution import ClaimedExecution
 from steady_queue.processes.concurrent import AtomicInteger
 
@@ -24,7 +25,8 @@ def post(self, execution: ClaimedExecution):
 
         def wrapped_execution():
             try:
-                execution.perform()
+                with AppExecutor.wrap_in_app_executor():
+                    execution.perform()
             finally:
                 self.available_threads.increment()
                 with self.mutex:

steady_queue/processes/registrable.py

@@ -2,6 +2,7 @@
 from typing import Optional
 
 import steady_queue
+from steady_queue.app_executor import AppExecutor
 from steady_queue.models import Process
 from steady_queue.processes.base import Base
 from steady_queue.processes.timer import TimerTask
@@ -67,12 +68,13 @@ def stop_heartbeat(self):
         logger.debug("stopped heartbeat for %s", self.name)
 
     def heartbeat(self):
-        try:
-            logger.debug("heartbeat from %s", self.name)
-            self.process.heartbeat()
-        except Process.DoesNotExist:
-            self.process = None
-            self.wake_up()
+        with AppExecutor.wrap_in_app_executor():
+            try:
+                logger.debug("heartbeat from %s", self.name)
+                self.process.heartbeat()
+            except Process.DoesNotExist:
+                self.process = None
+                self.wake_up()
 
     @property
     def process_id(self):

steady_queue/processes/runnable.py

@@ -1,7 +1,6 @@
 import logging
 import os
 
-from django.db import connections
 from steady_queue.processes.supervised import Supervised
 
 logger = logging.getLogger("steady_queue")
@@ -57,9 +56,6 @@ def is_all_work_completed(self) -> bool:
     def set_procline(self):
         pass
 
-    def reset_database_connections(self):
-        connections.close_all()
-
     @property
     def is_running_inline(self) -> bool:
         return self.mode == "inline"

steady_queue/processes/signals.py

@@ -41,7 +41,7 @@ def handle_signal(self, sig):
         if sig in (signal.SIGTERM, signal.SIGINT):
             self.stop()
             self.terminate_gracefully()
-        elif sig == self.SIGQUIT:
+        elif sig == signal.SIGQUIT:
             self.stop()
             self.terminate_immediately()
         else: