From d76b600a94db206caf661e1425bb72494d9f67ff Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Tue, 27 Jan 2026 09:38:25 -0800
Subject: [PATCH] Add flags to skip potentially problematic queries

---
 .../commands/document_perf_benchmark.py       | 33 +++++++++++++++----
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/documents/management/commands/document_perf_benchmark.py b/src/documents/management/commands/document_perf_benchmark.py
index 747bf776d..42f2e7f26 100644
--- a/src/documents/management/commands/document_perf_benchmark.py
+++ b/src/documents/management/commands/document_perf_benchmark.py
@@ -109,6 +109,16 @@ class Command(BaseCommand):
             default=1,
             help="How many custom field instances per document (default: 1)",
         )
+        parser.add_argument(
+            "--skip-tags",
+            action="store_true",
+            help="Skip tag document_count benchmarks (useful for large datasets on Postgres)",
+        )
+        parser.add_argument(
+            "--skip-custom-fields",
+            action="store_true",
+            help="Skip custom field document_count benchmarks",
+        )
         parser.add_argument(
             "--reuse-existing",
             action="store_true",
@@ -121,6 +131,9 @@ class Command(BaseCommand):
         )
 
     def handle(self, *args, **options):
+        # keep options for downstream checks
+        self.options = options
+
         document_total = options["documents"]
         owner_ratio = options["owner_ratio"]
         unowned_ratio = options["unowned_ratio"]
@@ -457,13 +470,19 @@ class Command(BaseCommand):
             iterations=iterations,
             fn=lambda: Document.objects.count(),
         )
-        self._time_tag_counts(iterations=iterations, prefix=prefix, user=target_user)
-        self._time_custom_field_counts(
-            iterations=iterations,
-            prefix=prefix,
-            user=target_user,
-            superuser=superuser,
-        )
+        if not self.options.get("skip_tags"):
+            self._time_tag_counts(
+                iterations=iterations,
+                prefix=prefix,
+                user=target_user,
+            )
+        if not self.options.get("skip_custom_fields"):
+            self._time_custom_field_counts(
+                iterations=iterations,
+                prefix=prefix,
+                user=target_user,
+                superuser=superuser,
+            )
 
     def _count_with_values_list(self, user) -> int:
         qs = get_objects_for_user_owner_aware(