feat(occurrences on eap): Implement double reads from EAP for organization traces API endpoint

src/sentry/api/endpoints/organization_traces.py

@@ -45,12 +45,15 @@
 from sentry.exceptions import InvalidSearchQuery
 from sentry.models.organization import Organization
 from sentry.models.project import Project
+from sentry.search.eap.occurrences.common_queries import count_occurrences_grouped_by_trace_ids
+from sentry.search.eap.occurrences.rollout_utils import EAPOccurrencesComparator
 from sentry.search.eap.types import SearchResolverConfig
 from sentry.search.events.builder.base import BaseQueryBuilder
 from sentry.search.events.builder.discover import DiscoverQueryBuilder
 from sentry.search.events.constants import TIMEOUT_SPAN_ERROR_MESSAGE
 from sentry.search.events.types import QueryBuilderConfig, SnubaParams, WhereType
 from sentry.snuba.dataset import Dataset
+from sentry.snuba.occurrences_rpc import OccurrenceCategory
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.spans_rpc import Spans
 from sentry.utils.numbers import clip
@@ -68,6 +71,13 @@ def is_trace_name_candidate(span):
     return span["span.op"] in CANDIDATE_SPAN_OPS
 
 
+def _reasonable_trace_count_map_match(snuba_map: dict[str, int], eap_map: dict[str, int]) -> bool:
+    if not set(eap_map).issubset(set(snuba_map)):
+        return False
+
+    return all(eap_map[trace_id] <= snuba_map[trace_id] for trace_id in eap_map)
+
+
 class TraceInterval(TypedDict):
     project: str | None
     sdkName: str | None
@@ -342,6 +352,49 @@ def enrich_eap_traces_with_extra_data(
             row["trace"]: row["count()"] for row in extra_results[1]["data"]
         }
 
+        organization_id = snuba_params.organization.id if snuba_params.organization else None
+        debug_context = {
+            "trace_ids": trace_ids,
+            "organization_id": organization_id,
+            "project_ids": [project.id for project in snuba_params.projects],
+            "start": snuba_params.start.isoformat() if snuba_params.start else None,
+            "end": snuba_params.end.isoformat() if snuba_params.end else None,
+        }
+
+        errors_callsite = "api.trace_explorer.traces_errors"
+        if EAPOccurrencesComparator.should_check_experiment(errors_callsite):
+            eap_traces_errors = count_occurrences_grouped_by_trace_ids(
+                snuba_params=snuba_params,
+                trace_ids=trace_ids,
+                referrer=Referrer.API_TRACE_EXPLORER_TRACES_ERRORS.value,
+                occurrence_category=OccurrenceCategory.ERROR,
+            )
+            traces_errors = EAPOccurrencesComparator.check_and_choose(
+                traces_errors,
+                eap_traces_errors,
+                errors_callsite,
+                is_experimental_data_a_null_result=len(eap_traces_errors) == 0,
+                reasonable_match_comparator=_reasonable_trace_count_map_match,
+                debug_context=debug_context,
+            )
+
+        occurrences_callsite = "api.trace_explorer.traces_occurrences"
+        if EAPOccurrencesComparator.should_check_experiment(occurrences_callsite):
+            eap_traces_occurrences = count_occurrences_grouped_by_trace_ids(
+                snuba_params=snuba_params,
+                trace_ids=trace_ids,
+                referrer=Referrer.API_TRACE_EXPLORER_TRACES_OCCURRENCES.value,
+                occurrence_category=OccurrenceCategory.GENERIC,
+            )
+            traces_occurrences = EAPOccurrencesComparator.check_and_choose(
+                traces_occurrences,
+                eap_traces_occurrences,
+                occurrences_callsite,
+                is_experimental_data_a_null_result=len(eap_traces_occurrences) == 0,
+                reasonable_match_comparator=_reasonable_trace_count_map_match,
+                debug_context=debug_context,
+            )
+
         self.enrich_traces_with_extra_data(
             traces,
             spans,

src/sentry/search/eap/occurrences/common_queries.py

@@ -75,3 +75,55 @@ def count_occurrences(
             },
         )
         return 0
+
+
+def count_occurrences_grouped_by_trace_ids(
+    snuba_params: SnubaParams,
+    trace_ids: Sequence[str],
+    referrer: str,
+    occurrence_category: OccurrenceCategory | None = None,
+) -> dict[str, int]:
+    """
+    Count occurrences grouped by trace IDs.
+
+    Returns:
+        A mapping of trace ID to occurrence count, or an empty dict if the query fails.
+    """
+    if not trace_ids:
+        return {}
+
+    query_string = f"trace:[{','.join(trace_ids)}]"
+
+    try:
+        result = Occurrences.run_table_query(
+            params=snuba_params,
+            query_string=query_string,
+            selected_columns=["trace", "count()"],
+            orderby=None,
+            offset=0,
+            limit=len(trace_ids),
+            referrer=referrer,
+            config=SearchResolverConfig(),
+            occurrence_category=occurrence_category,
+        )
+
+        return {
+            str(row["trace"]): int(row["count()"])
+            for row in result.get("data", [])
+            if row.get("trace") is not None
+        }
+    except Exception:
+        logger.exception(
+            "Fetching grouped trace occurrence counts from EAP failed",
+            extra={
+                "organization_id": (
+                    snuba_params.organization.id if snuba_params.organization else None
+                ),
+                "project_ids": [p.id for p in snuba_params.projects],
+                "trace_ids": list(trace_ids),
+                "start": snuba_params.start.isoformat() if snuba_params.start else None,
+                "end": snuba_params.end.isoformat() if snuba_params.end else None,
+                "referrer": referrer,
+            },
+        )
+        return {}

tests/snuba/search/test_eap_occurrences.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta
+from uuid import uuid4
 
 import pytest
 
-from sentry.search.eap.occurrences.common_queries import count_occurrences
+from sentry.search.eap.occurrences.common_queries import (
+    count_occurrences,
+    count_occurrences_grouped_by_trace_ids,
+)
 from sentry.search.eap.types import EAPResponse, SearchResolverConfig
 from sentry.search.events.types import SnubaParams
 from sentry.snuba.occurrences_rpc import OccurrenceCategory, Occurrences
@@ -241,3 +245,79 @@ def test_filters_by_occurrence_category(self) -> None:
 
         assert error_count == 3
         assert generic_count == 2
+
+    def test_counts_grouped_by_trace_ids(self) -> None:
+        group = self.create_group(project=self.project)
+        trace_id_1 = uuid4().hex
+        trace_id_2 = uuid4().hex
+        self.store_occurrences(
+            [
+                self.create_eap_occurrence(group_id=group.id, trace_id=trace_id_1),
+                self.create_eap_occurrence(group_id=group.id, trace_id=trace_id_1),
+                self.create_eap_occurrence(group_id=group.id, trace_id=trace_id_2),
+            ]
+        )
+
+        grouped = count_occurrences_grouped_by_trace_ids(
+            snuba_params=SnubaParams(
+                start=self.start,
+                end=self.end,
+                organization=self.organization,
+                projects=[self.project],
+            ),
+            trace_ids=[trace_id_1, trace_id_2],
+            referrer=self.referrer,
+        )
+        assert grouped == {trace_id_1: 2, trace_id_2: 1}
+
+    def test_counts_grouped_by_trace_ids_with_occurrence_category(self) -> None:
+        group = self.create_group(project=self.project)
+        trace_id = uuid4().hex
+        self.store_occurrences(
+            [
+                self.create_eap_occurrence(
+                    group_id=group.id, trace_id=trace_id, occurrence_type="error"
+                ),
+                self.create_eap_occurrence(
+                    group_id=group.id, trace_id=trace_id, occurrence_type="error"
+                ),
+                self.create_eap_occurrence(
+                    group_id=group.id, trace_id=trace_id, occurrence_type="generic"
+                ),
+            ]
+        )
+
+        snuba_params = SnubaParams(
+            start=self.start,
+            end=self.end,
+            organization=self.organization,
+            projects=[self.project],
+        )
+        grouped_errors = count_occurrences_grouped_by_trace_ids(
+            snuba_params=snuba_params,
+            trace_ids=[trace_id],
+            referrer=self.referrer,
+            occurrence_category=OccurrenceCategory.ERROR,
+        )
+        grouped_generic = count_occurrences_grouped_by_trace_ids(
+            snuba_params=snuba_params,
+            trace_ids=[trace_id],
+            referrer=self.referrer,
+            occurrence_category=OccurrenceCategory.GENERIC,
+        )
+
+        assert grouped_errors == {trace_id: 2}
+        assert grouped_generic == {trace_id: 1}
+
+    def test_counts_grouped_by_trace_ids_empty_trace_ids(self) -> None:
+        grouped = count_occurrences_grouped_by_trace_ids(
+            snuba_params=SnubaParams(
+                start=self.start,
+                end=self.end,
+                organization=self.organization,
+                projects=[self.project],
+            ),
+            trace_ids=[],
+            referrer=self.referrer,
+        )
+        assert grouped == {}