Handcrafts SQL queries a little more to reduce the query count and/or the amount of returned data (#6489)

src/documents/bulk_edit.py

@@ -24,12 +24,17 @@ from documents.tasks import update_document_archive_file
 logger = logging.getLogger("paperless.bulk_edit")
 
 
-def set_correspondent(doc_ids, correspondent):
-    if correspondent:
-        correspondent = Correspondent.objects.get(id=correspondent)
+def set_correspondent(doc_ids: list[int], correspondent):
 
-    qs = Document.objects.filter(Q(id__in=doc_ids) & ~Q(correspondent=correspondent))
-    affected_docs = [doc.id for doc in qs]
+    if correspondent:
+        correspondent = Correspondent.objects.only("pk").get(id=correspondent)
+
+    qs = (
+        Document.objects.filter(Q(id__in=doc_ids) & ~Q(correspondent=correspondent))
+        .select_related("correspondent")
+        .only("pk", "correspondent__id")
+    )
+    affected_docs = list(qs.values_list("pk", flat=True))
     qs.update(correspondent=correspondent)
 
     bulk_update_documents.delay(document_ids=affected_docs)
@@ -37,14 +42,18 @@ def set_correspondent(doc_ids, correspondent):
     return "OK"
 
 
-def set_storage_path(doc_ids, storage_path):
+def set_storage_path(doc_ids: list[int], storage_path):
     if storage_path:
-        storage_path = StoragePath.objects.get(id=storage_path)
+        storage_path = StoragePath.objects.only("pk").get(id=storage_path)
 
-    qs = Document.objects.filter(
-        Q(id__in=doc_ids) & ~Q(storage_path=storage_path),
+    qs = (
+        Document.objects.filter(
+            Q(id__in=doc_ids) & ~Q(storage_path=storage_path),
+        )
+        .select_related("storage_path")
+        .only("pk", "storage_path__id")
     )
-    affected_docs = [doc.id for doc in qs]
+    affected_docs = list(qs.values_list("pk", flat=True))
     qs.update(storage_path=storage_path)
 
     bulk_update_documents.delay(
@@ -54,12 +63,16 @@ def set_storage_path(doc_ids, storage_path):
     return "OK"
 
 
-def set_document_type(doc_ids, document_type):
+def set_document_type(doc_ids: list[int], document_type):
     if document_type:
-        document_type = DocumentType.objects.get(id=document_type)
+        document_type = DocumentType.objects.only("pk").get(id=document_type)
 
-    qs = Document.objects.filter(Q(id__in=doc_ids) & ~Q(document_type=document_type))
-    affected_docs = [doc.id for doc in qs]
+    qs = (
+        Document.objects.filter(Q(id__in=doc_ids) & ~Q(document_type=document_type))
+        .select_related("document_type")
+        .only("pk", "document_type__id")
+    )
+    affected_docs = list(qs.values_list("pk", flat=True))
     qs.update(document_type=document_type)
 
     bulk_update_documents.delay(document_ids=affected_docs)
@@ -67,9 +80,10 @@ def set_document_type(doc_ids, document_type):
     return "OK"
 
 
-def add_tag(doc_ids, tag):
-    qs = Document.objects.filter(Q(id__in=doc_ids) & ~Q(tags__id=tag))
-    affected_docs = [doc.id for doc in qs]
+def add_tag(doc_ids: list[int], tag: int):
+
+    qs = Document.objects.filter(Q(id__in=doc_ids) & ~Q(tags__id=tag)).only("pk")
+    affected_docs = list(qs.values_list("pk", flat=True))
 
     DocumentTagRelationship = Document.tags.through
 
@@ -82,9 +96,10 @@ def add_tag(doc_ids, tag):
     return "OK"
 
 
-def remove_tag(doc_ids, tag):
-    qs = Document.objects.filter(Q(id__in=doc_ids) & Q(tags__id=tag))
-    affected_docs = [doc.id for doc in qs]
+def remove_tag(doc_ids: list[int], tag: int):
+
+    qs = Document.objects.filter(Q(id__in=doc_ids) & Q(tags__id=tag)).only("pk")
+    affected_docs = list(qs.values_list("pk", flat=True))
 
     DocumentTagRelationship = Document.tags.through
 
@@ -97,9 +112,9 @@ def remove_tag(doc_ids, tag):
     return "OK"
 
 
-def modify_tags(doc_ids, add_tags, remove_tags):
-    qs = Document.objects.filter(id__in=doc_ids)
-    affected_docs = [doc.id for doc in qs]
+def modify_tags(doc_ids: list[int], add_tags: list[int], remove_tags: list[int]):
+    qs = Document.objects.filter(id__in=doc_ids).only("pk")
+    affected_docs = list(qs.values_list("pk", flat=True))
 
     DocumentTagRelationship = Document.tags.through
 
@@ -121,9 +136,9 @@ def modify_tags(doc_ids, add_tags, remove_tags):
     return "OK"
 
 
-def modify_custom_fields(doc_ids, add_custom_fields, remove_custom_fields):
-    qs = Document.objects.filter(id__in=doc_ids)
-    affected_docs = [doc.id for doc in qs]
+def modify_custom_fields(doc_ids: list[int], add_custom_fields, remove_custom_fields):
+    qs = Document.objects.filter(id__in=doc_ids).only("pk")
+    affected_docs = list(qs.values_list("pk", flat=True))
 
     fields_to_add = []
     for field in add_custom_fields:
@@ -145,7 +160,7 @@ def modify_custom_fields(doc_ids, add_custom_fields, remove_custom_fields):
     return "OK"
 
 
-def delete(doc_ids):
+def delete(doc_ids: list[int]):
     Document.objects.filter(id__in=doc_ids).delete()
 
     from documents import index
@@ -157,7 +172,7 @@ def delete(doc_ids):
     return "OK"
 
 
-def redo_ocr(doc_ids):
+def redo_ocr(doc_ids: list[int]):
     for document_id in doc_ids:
         update_document_archive_file.delay(
             document_id=document_id,
@@ -166,8 +181,8 @@ def redo_ocr(doc_ids):
     return "OK"
 
 
-def set_permissions(doc_ids, set_permissions, owner=None, merge=False):
-    qs = Document.objects.filter(id__in=doc_ids)
+def set_permissions(doc_ids: list[int], set_permissions, owner=None, merge=False):
+    qs = Document.objects.filter(id__in=doc_ids).select_related("owner")
 
     if merge:
         # If merging, only set owner for documents that don't have an owner
@@ -178,7 +193,7 @@ def set_permissions(doc_ids, set_permissions, owner=None, merge=False):
     for doc in qs:
         set_permissions_for_object(permissions=set_permissions, object=doc, merge=merge)
 
-    affected_docs = [doc.id for doc in qs]
+    affected_docs = list(qs.values_list("pk", flat=True))
 
     bulk_update_documents.delay(document_ids=affected_docs)