diff --git a/apps/worker/services/test_analytics/ta_process_flakes.py b/apps/worker/services/test_analytics/ta_process_flakes.py index 61cf26df09..8d244fe46a 100644 --- a/apps/worker/services/test_analytics/ta_process_flakes.py +++ b/apps/worker/services/test_analytics/ta_process_flakes.py @@ -33,12 +33,11 @@ def fetch_current_flakes(repo_id: int) -> dict[bytes, Flake]: } -def get_testruns(upload: ReportSession) -> QuerySet[Testrun]: - upload_filter = Q(upload_id=upload.id) - +def get_testruns_for_uploads(upload_ids: list[int]) -> QuerySet[Testrun]: # we won't process flakes for testruns older than 1 day return Testrun.objects.filter( - Q(timestamp__gte=timezone.now() - timedelta(days=1)) & upload_filter + Q(timestamp__gte=timezone.now() - timedelta(days=1)) + & Q(upload_id__in=upload_ids) ).order_by("timestamp") @@ -81,10 +80,8 @@ def handle_failure( @sentry_sdk.trace def process_single_upload( - upload: ReportSession, curr_flakes: dict[bytes, Flake], repo_id: int + testruns: list[Testrun], curr_flakes: dict[bytes, Flake], repo_id: int ): - testruns = get_testruns(upload) - for testrun in testruns: test_id = bytes(testrun.test_id) match testrun.outcome: @@ -98,19 +95,18 @@ def process_single_upload( case _: continue - Testrun.objects.bulk_update(testruns, ["outcome"]) - @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)) + upload_ids = [upload.id for upload in uploads] log.info( "process_flakes_for_commit: fetched uploads", - extra={"uploads": [upload.id for upload in uploads]}, + extra={"uploads": upload_ids}, ) curr_flakes = fetch_current_flakes(repo_id) @@ -120,13 +116,23 @@ def process_flakes_for_commit(repo_id: int, commit_id: str): extra={"flakes": [flake.test_id.hex() for flake in curr_flakes.values()]}, ) + # Fetch all testruns for all uploads in a single query, then group by upload_id + all_testruns = list(get_testruns_for_uploads(upload_ids)) + testruns_by_upload: dict[int, list[Testrun]] = {uid: [] for uid in upload_ids} + for testrun in all_testruns: + if testrun.upload_id in testruns_by_upload: + testruns_by_upload[testrun.upload_id].append(testrun) + for upload in uploads: - process_single_upload(upload, curr_flakes, repo_id) + process_single_upload(testruns_by_upload[upload.id], curr_flakes, repo_id) log.info( "process_flakes_for_commit: processed upload", extra={"upload": upload.id}, ) + # Bulk-update all testruns whose outcome may have been changed to "flaky_fail" + 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()]},