fix(worker): Resolve N+1 query in ProcessFlakesTask

apps/worker/services/test_analytics/ta_process_flakes.py

@@ -42,6 +42,23 @@ def get_testruns(upload: ReportSession) -> QuerySet[Testrun]:
     ).order_by("timestamp")
 
 
+def get_testruns_for_uploads(
+    uploads: list[ReportSession],
+) -> dict[int, list[Testrun]]:
+    """Batch-fetch all testruns for a list of uploads in a single query,
+    grouped by upload_id to avoid N+1 queries."""
+    upload_ids = [upload.id for upload in uploads]
+    testruns = Testrun.objects.filter(
+        upload_id__in=upload_ids,
+        timestamp__gte=timezone.now() - timedelta(days=1),
+    ).order_by("timestamp")
+
+    grouped: dict[int, list[Testrun]] = {upload_id: [] for upload_id in upload_ids}
+    for testrun in testruns:
+        grouped[testrun.upload_id].append(testrun)
+    return grouped
+
+
 def handle_pass(curr_flakes: dict[bytes, Flake], test_id: bytes):
     # possible that we expire it and stop caring about it
     if test_id not in curr_flakes:
@@ -81,9 +98,13 @@ def handle_failure(
 
 @sentry_sdk.trace
 def process_single_upload(
-    upload: ReportSession, curr_flakes: dict[bytes, Flake], repo_id: int
+    upload: ReportSession,
+    curr_flakes: dict[bytes, Flake],
+    repo_id: int,
+    testruns: list[Testrun] | None = None,
 ):
-    testruns = get_testruns(upload)
+    if testruns is None:
+        testruns = list(get_testruns(upload))
 
     for testrun in testruns:
         test_id = bytes(testrun.test_id)
@@ -98,15 +119,15 @@ def process_single_upload(
             case _:
                 continue
 
-    Testrun.objects.bulk_update(testruns, ["outcome"])
+    return testruns
 
 
 @sentry_sdk.trace
 def process_flakes_for_commit(repo_id: int, commit_id: str):
     log.info(
         "process_flakes_for_commit: starting processing",
     )
-    uploads = get_relevant_uploads(repo_id, commit_id)
+    uploads = list(get_relevant_uploads(repo_id, commit_id))
 
     log.info(
         "process_flakes_for_commit: fetched uploads",
@@ -120,13 +141,21 @@ def process_flakes_for_commit(repo_id: int, commit_id: str):
         extra={"flakes": [flake.test_id.hex() for flake in curr_flakes.values()]},
     )
 
+    testruns_by_upload = get_testruns_for_uploads(uploads)
+    all_testruns: list[Testrun] = []
+
     for upload in uploads:
-        process_single_upload(upload, curr_flakes, repo_id)
+        testruns = process_single_upload(
+            upload, curr_flakes, repo_id, testruns=testruns_by_upload[upload.id]
+        )
+        all_testruns.extend(testruns)
         log.info(
             "process_flakes_for_commit: processed upload",
             extra={"upload": upload.id},
         )
 
+    Testrun.objects.bulk_update(all_testruns, ["outcome"])
+
     log.info(
         "process_flakes_for_commit: bulk creating flakes",
         extra={"flakes": [flake.test_id.hex() for flake in curr_flakes.values()]},