Preload the {coverage,bundle}Status for commits (#1270)

Author: Swatinem

graphql_api/actions/commits.py

@@ -1,7 +1,9 @@
+from collections import defaultdict
 from typing import Optional
 
-from django.db.models import Prefetch, Q, QuerySet
+from django.db.models import Count, Prefetch, Q, QuerySet
 from django.db.models.functions import Lower, Substr
+from graphql import GraphQLResolveInfo
 
 from core.models import Commit, Pull, Repository
 from graphql_api.types.enums import CommitStatus
@@ -23,35 +25,51 @@ def pull_commits(pull: Pull) -> QuerySet[Commit]:
     return Commit.objects.filter(id__in=subquery).defer("_report")
 
 
+def load_commit_statuses(
+    commit_ids: list[int],
+) -> dict[int, dict[CommitReport.ReportType, str]]:
+    qs = (
+        CommitReport.objects.filter(commit__in=commit_ids)
+        .values_list("commit_id", "report_type", "sessions__state")
+        .annotate(sessions_count=Count("sessions"))
+    )
+
+    grouped: dict[tuple[int, CommitReport.ReportType], dict[str, int]] = defaultdict(
+        dict
+    )
+    for id, report_type, state, count in qs:
+        # The above query generates a `LEFT OUTER JOIN` with a proper `GROUP BY`.
+        # However, it is also yielding rows with a `NULL` state in case a report does not have any uploads.
+        if not report_type or not state:
+            continue
+        grouped[(id, report_type)][state] = count
+
+    results: dict[int, dict[CommitReport.ReportType, str]] = {
+        id: {} for id in commit_ids
+    }
+    for (id, report_type), states in grouped.items():
+        status = CommitStatus.COMPLETED.value
+        if states.get("error", 0) > 0:
+            status = CommitStatus.ERROR.value
+        elif states.get("uploaded", 0) > 0:
+            status = CommitStatus.PENDING.value
+
+        results[id][report_type] = status
+
+    return results
+
+
 def commit_status(
-    commit: Commit, report_type: CommitReport.ReportType
-) -> Optional[CommitStatus]:
-    report = CommitReport.objects.filter(report_type=report_type, commit=commit).first()
-    if not report:
-        return None
-
-    sessions = report.sessions.all()
-    if not sessions:
-        return None
-
-    # Only care about these 3 states, ignoring fully and partially overwritten
-    upload_states = [
-        s.state for s in sessions if s.state in ["processed", "uploaded", "error"]
-    ]
-
-    has_error, has_pending = False, False
-    for state in upload_states:
-        if state == "error":
-            has_error = True
-        if state == "uploaded":
-            has_pending = True
-
-    # Prioritize returning error over pending
-    if has_error:
-        return CommitStatus.ERROR.value
-    if has_pending:
-        return CommitStatus.PENDING.value
-    return CommitStatus.COMPLETED.value
+    info: GraphQLResolveInfo, commit: Commit, report_type: CommitReport.ReportType
+) -> str | None:
+    commit_statuses = info.context.setdefault("commit_statuses", {})
+    commit_status = commit_statuses.get(commit.id)
+    if commit_status is None:
+        updated_statuses = load_commit_statuses([commit.id])
+        commit_statuses.update(updated_statuses)
+        commit_status = updated_statuses[commit.id]
+
+    return commit_status.get(report_type)
 
 
 def repo_commits(
@@ -100,11 +118,17 @@ def repo_commits(
     coverage_status = filters.get("coverage_status")
 
     if coverage_status:
+        # FIXME(swatinem):
+        # This filter here is insane, it resolves *all* the results in the unbounded queryset,
+        # just to check the status, and to then add it as another restricting filter.
+        # I’m pretty sure this will completely break the server if anyone actually uses this filter, lol.
+        commit_ids = [commit.id for commit in queryset]
+        commit_statuses = load_commit_statuses(commit_ids)
+
         to_be_included = [
-            commit.id
-            for commit in queryset
-            if commit_status(commit, CommitReport.ReportType.COVERAGE)
-            in coverage_status
+            id
+            for id, statuses in commit_statuses.items()
+            if statuses.get(CommitReport.ReportType.COVERAGE) in coverage_status
         ]
         queryset = queryset.filter(id__in=to_be_included)
 

graphql_api/types/commit/commit.py

@@ -36,7 +36,6 @@
     MissingHeadReport,
 )
 from graphql_api.types.enums import (
-    CommitStatus,
     OrderingDirection,
     PathContentDisplayType,
 )
@@ -304,15 +303,15 @@ async def resolve_total_uploads(commit, info):
 @commit_bindable.field("bundleStatus")
 @sync_to_async
 @sentry_sdk.trace
-def resolve_bundle_status(commit: Commit, info) -> Optional[CommitStatus]:
-    return commit_status(commit, CommitReport.ReportType.BUNDLE_ANALYSIS)
+def resolve_bundle_status(commit: Commit, info: GraphQLResolveInfo) -> str | None:
+    return commit_status(info, commit, CommitReport.ReportType.BUNDLE_ANALYSIS)
 
 
 @commit_bindable.field("coverageStatus")
 @sync_to_async
 @sentry_sdk.trace
-def resolve_coverage_status(commit: Commit, info) -> Optional[CommitStatus]:
-    return commit_status(commit, CommitReport.ReportType.COVERAGE)
+def resolve_coverage_status(commit: Commit, info: GraphQLResolveInfo) -> str | None:
+    return commit_status(info, commit, CommitReport.ReportType.COVERAGE)
 
 
 @commit_bindable.field("coverageAnalytics")

graphql_api/types/repository/repository.py

@@ -13,10 +13,11 @@
 
 from codecov_auth.models import SERVICE_GITHUB, SERVICE_GITHUB_ENTERPRISE, Owner
 from core.models import Branch, Commit, Pull, Repository
-from graphql_api.actions.commits import repo_commits
+from graphql_api.actions.commits import load_commit_statuses, repo_commits
 from graphql_api.dataloader.commit import CommitLoader
 from graphql_api.dataloader.owner import OwnerLoader
 from graphql_api.helpers.connection import queryset_to_connection
+from graphql_api.helpers.requested_fields import selected_fields
 from graphql_api.types.coverage_analytics.coverage_analytics import (
     CoverageAnalyticsProps,
 )
@@ -123,6 +124,11 @@ async def resolve_pulls(
     )
 
 
+# the `requested_fields` here are prefixed with `edges.node`, as this is a `Connection`
+# and using `commits { edges { node { ... } } }` is the way this is queried.
+STATUS_FIELDS = {"edges.node.coverageStatus", "edges.node.bundleStatus"}
+
+
 @repository_bindable.field("commits")
 async def resolve_commits(
     repository: Repository,
@@ -144,6 +150,15 @@ async def resolve_commits(
         loader = CommitLoader.loader(info, repository.repoid)
         loader.cache(commit)
 
+    requested_fields = selected_fields(info)
+    should_load_statuses = not requested_fields.isdisjoint(STATUS_FIELDS)
+
+    if should_load_statuses:
+        commit_ids = [edge["node"].id for edge in connection.edges]
+        info.context["commit_statuses"] = await sync_to_async(load_commit_statuses)(
+            commit_ids
+        )
+
     return connection