perf: optimize import persistence for 0.3.14

scripts/benchmark.py

@@ -20,7 +20,6 @@ import django
 from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.db import connection
-from django.test.utils import CaptureQueriesContext
 from django.utils import timezone
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings")
@@ -29,7 +28,7 @@ django.setup()
 from apps.jobs.models import Employer, JobPosting, JobSourceAlias, ScoreRun  # noqa: E402
 from apps.jobs.services.dedupe import find_existing_job  # noqa: E402
 from apps.jobs.services.normalization import CanonicalJobDraft, normalize_token  # noqa: E402
-from apps.jobs.services.pipeline import process_raw_document  # noqa: E402
+from apps.jobs.services.pipeline import PersistenceContext, process_raw_document  # noqa: E402
 from apps.jobs.services.scoring import calculate_score, rescore_jobs_with_profiles  # noqa: E402
 from apps.profiles.models import SearchProfile  # noqa: E402
 from apps.sources.adapters.registry import registry  # noqa: E402
@@ -575,15 +574,25 @@ def _clean_benchmark(artifacts: BenchmarkArtifact) -> None:
         get_user_model().objects.filter(pk=artifacts.user_id).delete()
 
 
+class QueryCounter:
+    """Count database executions without Django's bounded debug-query log."""
+
+    def __init__(self) -> None:
+        self.count = 0
+        self.types: dict[str, int] = {}
+
+    def __call__(self, execute: Any, sql: str, params: Any, many: bool, context: Any) -> Any:
+        self.count += 1
+        query_type = sql.lstrip().split(None, 1)[0].upper() if sql.strip() else "UNKNOWN"
+        self.types[query_type] = self.types.get(query_type, 0) + 1
+        return execute(sql, params, many, context)
+
+
 @contextmanager
 def _query_capture() -> Any:
-    original_force_debug_cursor = connection.force_debug_cursor
-    connection.force_debug_cursor = True
-    try:
-        with CaptureQueriesContext(connection) as captured:
-            yield captured
-    finally:
-        connection.force_debug_cursor = original_force_debug_cursor
+    counter = QueryCounter()
+    with connection.execute_wrapper(counter):
+        yield counter
 
 
 def run_performance_benchmark(
@@ -619,6 +628,7 @@ def run_performance_benchmark(
     }
     tracemalloc.start()
     import_timings: list[float] = []
+    persistence_context = PersistenceContext()
     rescore_timings: list[float] = []
     list_timings: list[float] = []
     detail_timings: list[float] = []
@@ -646,9 +656,10 @@ def run_performance_benchmark(
                         retain_until=timezone.now() + timezone.timedelta(days=7),
                     )
                     start = perf_counter()
-                    process_raw_document(document)
+                    process_raw_document(document, context=persistence_context)
                     import_timings.append((perf_counter() - start) * 1000)
-        import_query_count = len(captured_import.captured_queries)
+        import_query_count = captured_import.count
+        import_query_types = captured_import.types
         query_count_total += import_query_count
 
         candidates = list(
@@ -666,7 +677,8 @@ def run_performance_benchmark(
                 batch_start = perf_counter()
                 rescore_jobs_with_profiles(batch, profile_id=profile.id)
                 rescore_timings.append((perf_counter() - batch_start) * 1000)
-        rescore_query_count = len(captured_rescore.captured_queries)
+        rescore_query_count = captured_rescore.count
+        rescore_query_types = captured_rescore.types
         query_count_total += rescore_query_count
 
         with _query_capture() as captured_list:
@@ -695,8 +707,10 @@ def run_performance_benchmark(
                 )
                 detail_timings.append((perf_counter() - start) * 1000)
 
-        list_query_count = len(captured_list.captured_queries)
-        detail_query_count = len(captured_detail.captured_queries)
+        list_query_count = captured_list.count
+        detail_query_count = captured_detail.count
+        list_query_types = captured_list.types
+        detail_query_types = captured_detail.types
         query_count_total += list_query_count + detail_query_count
 
         current_bytes, peak_bytes = tracemalloc.get_traced_memory()
@@ -740,6 +754,7 @@ def run_performance_benchmark(
             "p50_ms": timing_stats["import_p50_ms"],
             "p95_ms": timing_stats["import_p95_ms"],
             "query_count": import_query_count,
+            "query_types": import_query_types,
             "timing_count": len(import_timings),
         },
         "rescore": {
@@ -747,11 +762,14 @@ def run_performance_benchmark(
             "p50_ms": timing_stats["rescore_p50_ms"],
             "p95_ms": timing_stats["rescore_p95_ms"],
             "query_count": rescore_query_count,
+            "query_types": rescore_query_types,
             "timing_count": len(rescore_timings),
         },
         "dashboard": {
             "list_query_count": list_query_count,
             "detail_query_count": detail_query_count,
+            "list_query_types": list_query_types,
+            "detail_query_types": detail_query_types,
             "list_p95_ms": timing_stats["list_p95_ms"],
             "detail_p95_ms": timing_stats["detail_p95_ms"],
             "timing_count": len(list_timings) + len(detail_timings),