Fix audit task DB connection leaks in E2E runs

api/audit/tasks.py

@@ -2,12 +2,15 @@
 import typing
 from datetime import datetime
 
+from django.conf import settings
 from django.contrib.auth import get_user_model
+from django.db import connections
 from django.utils import timezone
 from task_processor.decorators import (
     register_task_handler,
 )
 from task_processor.models import TaskPriority
+from task_processor.task_run_method import TaskRunMethod
 
 from audit.constants import (
     FEATURE_STATE_UPDATED_BY_CHANGE_REQUEST_MESSAGE,
@@ -18,6 +21,11 @@
 logger = logging.getLogger(__name__)
 
 
+def _close_thread_database_connections() -> None:
+    if settings.TASK_RUN_METHOD == TaskRunMethod.SEPARATE_THREAD:
+        connections.close_all()
+
+
 @register_task_handler(priority=TaskPriority.HIGHEST)
 def create_feature_state_went_live_audit_log(feature_state_id: int):  # type: ignore[no-untyped-def]
     _create_feature_state_audit_log_for_change_request(
@@ -83,59 +91,64 @@ def create_audit_log_from_historical_record(  # type: ignore[no-untyped-def]
     history_user_id: typing.Optional[int],
     history_record_class_path: str,
 ):
-    model_class = AuditLog.get_history_record_model_class(history_record_class_path)
-    history_instance = model_class.objects.get(history_id=history_instance_id)  # type: ignore[attr-defined]
-
-    if (
-        history_instance.history_type == "~"
-        and history_instance.prev_record
-        and not history_instance.diff_against(history_instance.prev_record).changes
-    ):
-        return
-
-    instance = history_instance.instance
-    if instance.get_skip_create_audit_log():
-        return
-
-    if history_user_id is not None:
-        user_model = get_user_model()
-        history_user = user_model.objects.filter(id=history_user_id).first()
-    else:
-        history_user = instance.get_audit_log_author(history_instance)
-
-    if not (history_user or history_instance.master_api_key):
-        return
-
-    environment, project = instance.get_environment_and_project()
-
-    related_object_id = instance.get_audit_log_related_object_id(history_instance)
-    related_object_type = instance.get_audit_log_related_object_type(history_instance)
-
-    if not related_object_id:
-        return
-
-    log_message = {
-        "+": instance.get_create_log_message,
-        "-": instance.get_delete_log_message,
-        "~": instance.get_update_log_message,
-    }[history_instance.history_type](history_instance)
-
-    if not log_message:
-        return
+    try:
+        model_class = AuditLog.get_history_record_model_class(history_record_class_path)
+        history_instance = model_class.objects.get(history_id=history_instance_id)  # type: ignore[attr-defined]
+
+        if (
+            history_instance.history_type == "~"
+            and history_instance.prev_record
+            and not history_instance.diff_against(history_instance.prev_record).changes
+        ):
+            return
+
+        instance = history_instance.instance
+        if instance.get_skip_create_audit_log():
+            return
+
+        if history_user_id is not None:
+            user_model = get_user_model()
+            history_user = user_model.objects.filter(id=history_user_id).first()
+        else:
+            history_user = instance.get_audit_log_author(history_instance)
+
+        if not (history_user or history_instance.master_api_key):
+            return
+
+        environment, project = instance.get_environment_and_project()
+
+        related_object_id = instance.get_audit_log_related_object_id(history_instance)
+        related_object_type = instance.get_audit_log_related_object_type(
+            history_instance
+        )
 
-    AuditLog.objects.create(
-        history_record_id=history_instance.history_id,
-        history_record_class_path=history_record_class_path,
-        environment=environment,
-        project=project,
-        author=history_user,
-        related_object_id=related_object_id,
-        related_object_type=related_object_type.name,
-        log=log_message,
-        master_api_key=history_instance.master_api_key,
-        created_date=history_instance.history_date,
-        **instance.get_extra_audit_log_kwargs(history_instance),
-    )
+        if not related_object_id:
+            return
+
+        log_message = {
+            "+": instance.get_create_log_message,
+            "-": instance.get_delete_log_message,
+            "~": instance.get_update_log_message,
+        }[history_instance.history_type](history_instance)
+
+        if not log_message:
+            return
+
+        AuditLog.objects.create(
+            history_record_id=history_instance.history_id,
+            history_record_class_path=history_record_class_path,
+            environment=environment,
+            project=project,
+            author=history_user,
+            related_object_id=related_object_id,
+            related_object_type=related_object_type.name,
+            log=log_message,
+            master_api_key=history_instance.master_api_key,
+            created_date=history_instance.history_date,
+            **instance.get_extra_audit_log_kwargs(history_instance),
+        )
+    finally:
+        _close_thread_database_connections()
 
 
 @register_task_handler()

api/tests/unit/audit/test_unit_audit_tasks.py

@@ -158,6 +158,7 @@
 def test_create_audit_log_from_historical_record__valid_record__creates_audit_log_with_correct_fields(  # type: ignore[no-untyped-def]
     mocker,
     monkeypatch,
+    settings,
 ):
     # Given
     log_message = "a log message"
@@ -200,6 +201,7 @@
     mocked_audit_log_model_class.get_history_record_model_class.return_value = (
         mocked_historical_record_model_class
     )
+    mocked_close_all = mocker.patch("audit.tasks.connections.close_all")
 
     history_record_class_path = (
         f"app.models.{mocked_historical_record_model_class.name}"
@@ -223,6 +225,98 @@
         master_api_key=None,
         created_date=history_instance.history_date,
     )
+    mocked_close_all.assert_not_called()
+
+
+def test_create_audit_log_from_historical_record__separate_thread__closes_database_connections(
+    mocker,
+    settings,
+):
+    # Given
+    settings.TASK_RUN_METHOD = TaskRunMethod.SEPARATE_THREAD
+
+    log_message = "a log message"
+    related_object_id = 1
+    related_object_type = RelatedObjectType.ENVIRONMENT
+
+    mock_environment = mocker.MagicMock(spec=Environment)
+
+    instance = mocker.MagicMock()
+    instance.get_skip_create_audit_log.return_value = False
+    instance.get_audit_log_author.return_value = None
+    instance.get_create_log_message.return_value = log_message
+    instance.get_environment_and_project.return_value = mock_environment, None
+    instance.get_audit_log_related_object_id.return_value = related_object_id
+    instance.get_audit_log_related_object_type.return_value = related_object_type
+    instance.get_extra_audit_log_kwargs.return_value = {}
+    history_instance = mocker.MagicMock(
+        history_id=1,
+        instance=instance,
+        master_api_key=None,
+        history_type="+",
+        history_date=timezone.now(),
+    )
+
+    history_user = mocker.MagicMock()
+    history_user.id = 1
+
+    mocked_historical_record_model_class = mocker.MagicMock(
+        name="DummyHistoricalRecordModelClass"
+    )
+    mocked_historical_record_model_class.objects.get.return_value = history_instance
+
+    mocked_user_model_class = mocker.MagicMock()
+    mocker.patch("audit.tasks.get_user_model", return_value=mocked_user_model_class)
+    mocked_user_model_class.objects.filter.return_value.first.return_value = (
+        history_user
+    )
+
+    mocked_audit_log_model_class = mocker.patch("audit.tasks.AuditLog")
+    mocked_audit_log_model_class.get_history_record_model_class.return_value = (
+        mocked_historical_record_model_class
+    )
+    mocked_close_all = mocker.patch("audit.tasks.connections.close_all")
+
+    history_record_class_path = (
+        f"app.models.{mocked_historical_record_model_class.name}"
+    )
+
+    # When
+    create_audit_log_from_historical_record(
+        history_instance.history_id, history_user.id, history_record_class_path
+    )
+
+    # Then
+    mocked_close_all.assert_called_once_with()
+
+
+def test_create_audit_log_from_historical_record__separate_thread_exception__closes_database_connections(
+    mocker,
+    settings,
+) -> None:
+    # Given
+    settings.TASK_RUN_METHOD = TaskRunMethod.SEPARATE_THREAD
+
+    mocked_historical_record_model_class = mocker.MagicMock(
+        name="DummyHistoricalRecordModelClass"
+    )
+    mocked_historical_record_model_class.objects.get.side_effect = RuntimeError("boom")
+
+    mocked_audit_log_model_class = mocker.patch("audit.tasks.AuditLog")
+    mocked_audit_log_model_class.get_history_record_model_class.return_value = (
+        mocked_historical_record_model_class
+    )
+    mocked_close_all = mocker.patch("audit.tasks.connections.close_all")
+
+    history_record_class_path = (
+        f"app.models.{mocked_historical_record_model_class.name}"
+    )
+
+    # When / Then
+    with pytest.raises(RuntimeError, match="boom"):
+        create_audit_log_from_historical_record(1, None, history_record_class_path)
+
+    mocked_close_all.assert_called_once_with()
 
 
 def test_create_audit_log_from_historical_record__cascade_deleted_feature_segment__does_nothing(