From 15722229338e21c9a68324377c1e5c2ac0b9b934 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 18 Jun 2026 16:20:53 +0000 Subject: [PATCH 1/2] fix: close audit task db connections --- api/audit/tasks.py | 117 ++++++++++-------- api/tests/unit/audit/test_unit_audit_tasks.py | 96 ++++++++++++++ 2 files changed, 160 insertions(+), 53 deletions(-) diff --git a/api/audit/tasks.py b/api/audit/tasks.py index ca627e3b532c..6702f1dc26dd 100644 --- a/api/audit/tasks.py +++ b/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,62 @@ 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 - - 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), - ) + 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) + + 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() diff --git a/api/tests/unit/audit/test_unit_audit_tasks.py b/api/tests/unit/audit/test_unit_audit_tasks.py index 97192bb17ab4..193a0754fe7e 100644 --- a/api/tests/unit/audit/test_unit_audit_tasks.py +++ b/api/tests/unit/audit/test_unit_audit_tasks.py @@ -158,6 +158,7 @@ def test_create_audit_log_from_historical_record__skip_audit_log_true__does_not_ 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 @@ def test_create_audit_log_from_historical_record__valid_record__creates_audit_lo 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,100 @@ def test_create_audit_log_from_historical_record__valid_record__creates_audit_lo 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( From 0f8d845fc4021055659fe3b1ff4246cc3ca6b6fa Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 18 Jun 2026 16:41:37 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- api/audit/tasks.py | 4 +++- api/tests/unit/audit/test_unit_audit_tasks.py | 4 +--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/api/audit/tasks.py b/api/audit/tasks.py index 6702f1dc26dd..3e6fd232780e 100644 --- a/api/audit/tasks.py +++ b/api/audit/tasks.py @@ -118,7 +118,9 @@ def create_audit_log_from_historical_record( # type: ignore[no-untyped-def] 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) + related_object_type = instance.get_audit_log_related_object_type( + history_instance + ) if not related_object_id: return diff --git a/api/tests/unit/audit/test_unit_audit_tasks.py b/api/tests/unit/audit/test_unit_audit_tasks.py index 193a0754fe7e..cb954594181e 100644 --- a/api/tests/unit/audit/test_unit_audit_tasks.py +++ b/api/tests/unit/audit/test_unit_audit_tasks.py @@ -300,9 +300,7 @@ def test_create_audit_log_from_historical_record__separate_thread_exception__clo mocked_historical_record_model_class = mocker.MagicMock( name="DummyHistoricalRecordModelClass" ) - mocked_historical_record_model_class.objects.get.side_effect = RuntimeError( - "boom" - ) + 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 = (