Closing DB connections that are unused

dojo/celery.py

@@ -3,7 +3,7 @@
 from logging.config import dictConfig
 
 from celery import Celery, Task
-from celery.signals import setup_logging
+from celery.signals import setup_logging, task_postrun, task_prerun
 from django.conf import settings
 
 logger = logging.getLogger(__name__)
@@ -119,6 +119,22 @@ def config_loggers(*args, **kwags):
     dictConfig(settings.LOGGING)
 
 
+@task_prerun.connect
+def close_old_db_connections_before_task(task=None, **kwargs):
+    from django.db import close_old_connections  # noqa: PLC0415
+    if task is not None and getattr(task.request, "is_eager", False):
+        return
+    close_old_connections()
+
+
+@task_postrun.connect
+def close_old_db_connections_after_task(task=None, **kwargs):
+    from django.db import close_old_connections  # noqa: PLC0415
+    if task is not None and getattr(task.request, "is_eager", False):
+        return
+    close_old_connections()
+
+
 # from celery import current_app
 
 # _ = current_app.loader.import_default_modules()

dojo/settings/settings.dist.py

@@ -462,6 +462,9 @@ def generate_url(scheme, double_slashes, user, password, host, port, path, param
         },
     }
 
+DATABASES["default"]["CONN_MAX_AGE"] = 3600
+DATABASES["default"]["CONN_HEALTH_CHECKS"] = True
+
 # Track migrations through source control rather than making migrations locally
 if env("DD_TRACK_MIGRATIONS"):
     MIGRATION_MODULES = {"dojo": "dojo.db_migrations"}

unittests/test_celery_db_connections.py

@@ -0,0 +1,65 @@
+import logging
+from unittest.mock import Mock, patch
+
+from dojo.celery import (
+    close_old_db_connections_after_task,
+    close_old_db_connections_before_task,
+)
+
+from .dojo_test_case import DojoTestCase
+
+logger = logging.getLogger(__name__)
+
+
+class TestCloseOldDbConnectionsBeforeTask(DojoTestCase):
+
+    """Tests for the task_prerun signal handler that closes stale DB connections."""
+
+    @patch("django.db.close_old_connections")
+    def test_eager_task_skips_close(self, mock_close):
+        """Eager tasks (used in unit tests) should not close DB connections."""
+        task = Mock()
+        task.request.is_eager = True
+        close_old_db_connections_before_task(task=task)
+        mock_close.assert_not_called()
+
+    @patch("django.db.close_old_connections")
+    def test_non_eager_task_closes_connections(self, mock_close):
+        """Non-eager (normal async) tasks should close stale DB connections."""
+        task = Mock()
+        task.request.is_eager = False
+        close_old_db_connections_before_task(task=task)
+        mock_close.assert_called_once()
+
+    @patch("django.db.close_old_connections")
+    def test_none_task_closes_connections(self, mock_close):
+        """If task is None, connections should still be closed."""
+        close_old_db_connections_before_task(task=None)
+        mock_close.assert_called_once()
+
+
+class TestCloseOldDbConnectionsAfterTask(DojoTestCase):
+
+    """Tests for the task_postrun signal handler that closes stale DB connections."""
+
+    @patch("django.db.close_old_connections")
+    def test_eager_task_skips_close(self, mock_close):
+        """Eager tasks (used in unit tests) should not close DB connections."""
+        task = Mock()
+        task.request.is_eager = True
+        close_old_db_connections_after_task(task=task)
+        mock_close.assert_not_called()
+
+    @patch("django.db.close_old_connections")
+    def test_non_eager_task_closes_connections(self, mock_close):
+        """Non-eager (normal async) tasks should close stale DB connections."""
+        task = Mock()
+        task.request.is_eager = False
+        close_old_db_connections_after_task(task=task)
+        mock_close.assert_called_once()
+
+    @patch("django.db.close_old_connections")
+    def test_none_task_closes_connections(self, mock_close):
+        """If task is None, connections should still be closed."""
+        close_old_db_connections_after_task(task=None)
+        mock_close.assert_called_once()