From 612e6341a339cdb57a51f62485485dbad2718129 Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Mon, 25 Sep 2023 12:52:22 -0700
Subject: [PATCH] Ensures all old connections are closed in certain longer
 running locations (#4265)

---
 src/documents/management/commands/document_consumer.py | 2 ++
 src/documents/signals/handlers.py                      | 6 ++++++
 2 files changed, 8 insertions(+)

diff --git a/src/documents/management/commands/document_consumer.py b/src/documents/management/commands/document_consumer.py
index 0cc4dc3ef..64812ed88 100644
--- a/src/documents/management/commands/document_consumer.py
+++ b/src/documents/management/commands/document_consumer.py
@@ -9,6 +9,7 @@ from time import monotonic
 from time import sleep
 from typing import Final
 
+from django import db
 from django.conf import settings
 from django.core.management.base import BaseCommand
 from django.core.management.base import CommandError
@@ -38,6 +39,7 @@ def _tags_from_path(filepath) -> set[Tag]:
 
     Returns set of Tag models
     """
+    db.close_old_connections()
     tag_ids = set()
     path_parts = Path(filepath).relative_to(settings.CONSUMPTION_DIR).parent.parts
     for part in path_parts:
diff --git a/src/documents/signals/handlers.py b/src/documents/signals/handlers.py
index 4e0d13c20..117e3c38d 100644
--- a/src/documents/signals/handlers.py
+++ b/src/documents/signals/handlers.py
@@ -14,6 +14,7 @@ from django.contrib.admin.models import LogEntry
 from django.contrib.auth.models import User
 from django.contrib.contenttypes.models import ContentType
 from django.db import DatabaseError
+from django.db import close_old_connections
 from django.db import models
 from django.db.models import Q
 from django.dispatch import receiver
@@ -529,6 +530,8 @@ def before_task_publish_handler(sender=None, headers=None, body=None, **kwargs):
         return
 
     try:
+        close_old_connections()
+
         task_args = body[0]
         input_doc, _ = task_args
 
@@ -560,6 +563,7 @@ def task_prerun_handler(sender=None, task_id=None, task=None, **kwargs):
     https://docs.celeryq.dev/en/stable/userguide/signals.html#task-prerun
     """
     try:
+        close_old_connections()
         task_instance = PaperlessTask.objects.filter(task_id=task_id).first()
 
         if task_instance is not None:
@@ -587,6 +591,7 @@ def task_postrun_handler(
     https://docs.celeryq.dev/en/stable/userguide/signals.html#task-postrun
     """
     try:
+        close_old_connections()
         task_instance = PaperlessTask.objects.filter(task_id=task_id).first()
 
         if task_instance is not None:
@@ -615,6 +620,7 @@ def task_failure_handler(
     https://docs.celeryq.dev/en/stable/userguide/signals.html#task-failure
     """
     try:
+        close_old_connections()
         task_instance = PaperlessTask.objects.filter(task_id=task_id).first()
 
         if task_instance is not None and task_instance.result is None: