diff --git a/apps/worker/services/test_analytics/ta_process_flakes.py b/apps/worker/services/test_analytics/ta_process_flakes.py index 61cf26df09..567600259b 100644 --- a/apps/worker/services/test_analytics/ta_process_flakes.py +++ b/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,7 +119,7 @@ def process_single_upload( case _: continue - Testrun.objects.bulk_update(testruns, ["outcome"]) + return testruns @sentry_sdk.trace @@ -106,7 +127,7 @@ 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()]},