diff --git a/src/documents/signals/handlers.py b/src/documents/signals/handlers.py index fe1d1f94d..30c682cbf 100644 --- a/src/documents/signals/handlers.py +++ b/src/documents/signals/handlers.py @@ -1134,6 +1134,11 @@ def before_task_publish_handler( return try: + # Close stale connections without disrupting a transaction publishing a task + for connection in connections.all(initialized_only=True): + if not connection.in_atomic_block: + connection.close_if_unusable_or_obsolete() + _, task_kwargs, _ = body task_id = headers["id"] diff --git a/src/documents/tests/test_task_signals.py b/src/documents/tests/test_task_signals.py index 2c9fde16f..83652b5df 100644 --- a/src/documents/tests/test_task_signals.py +++ b/src/documents/tests/test_task_signals.py @@ -56,6 +56,32 @@ def send_publish( @pytest.mark.django_db class TestBeforeTaskPublishHandler: + @mock.patch("documents.signals.handlers.connections.all") + def test_closes_old_connections_outside_atomic_blocks( + self, + connections_all, + ) -> None: + connection = mock.Mock(in_atomic_block=False) + connections_all.return_value = [connection] + + task_id = send_publish("documents.tasks.train_classifier", (), {}) + + connection.close_if_unusable_or_obsolete.assert_called_once_with() + assert PaperlessTask.objects.filter(task_id=task_id).exists() + + @mock.patch("documents.signals.handlers.connections.all") + def test_keeps_connections_open_inside_atomic_blocks( + self, + connections_all, + ) -> None: + connection = mock.Mock(in_atomic_block=True) + connections_all.return_value = [connection] + + task_id = send_publish("documents.tasks.train_classifier", (), {}) + + connection.close_if_unusable_or_obsolete.assert_not_called() + assert PaperlessTask.objects.filter(task_id=task_id).exists() + def test_creates_task_for_consume_file( self, consume_input_doc,