fix(eap-span): n+1 on opportunity and perf score functions (#91808)

In both the `performance_score` and `opportunity_score` functions, we
were doing something like

```python
for vital in vitals
   vital_count = get_count_of_vital(vital) # this is a new query with each iteration
   if (vital_count > 0)
      # account for this vital in the score
```

This pr updates `get_count_of_vital` to query for all vitals at once,
and then cache the result. This way we don't have to make so many
queries and just reuse the result of any subsequent vital

The cache is stored on the `resolver` instance, therefore it is only
applied per request and won't bleed into other requests.

Author: DominikB2014

src/sentry/search/eap/columns.py

@@ -20,6 +20,7 @@
 
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap import constants
+from sentry.search.eap.types import EAPResponse
 from sentry.search.events.types import SnubaParams
 
 ResolvedArgument: TypeAlias = AttributeKey | str | int | float
@@ -29,6 +30,7 @@
 class ResolverSettings(TypedDict):
     extrapolation_mode: ExtrapolationMode.ValueType
     snuba_params: SnubaParams
+    query_result_cache: dict[str, EAPResponse]
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -259,6 +261,7 @@ def resolve(
         search_type: constants.SearchType,
         resolved_arguments: ResolvedArguments,
         snuba_params: SnubaParams,
+        query_result_cache: dict[str, EAPResponse],
     ) -> ResolvedFormula | ResolvedAggregate | ResolvedConditionalAggregate:
         raise NotImplementedError()
 
@@ -278,6 +281,7 @@ def resolve(
         search_type: constants.SearchType,
         resolved_arguments: ResolvedArguments,
         snuba_params: SnubaParams,
+        query_result_cache: dict[str, EAPResponse],
     ) -> ResolvedAggregate:
         if len(resolved_arguments) > 1:
             raise InvalidSearchQuery(
@@ -324,6 +328,7 @@ def resolve(
         search_type: constants.SearchType,
         resolved_arguments: ResolvedArguments,
         snuba_params: SnubaParams,
+        query_result_cache: dict[str, EAPResponse],
     ) -> ResolvedConditionalAggregate:
         key, filter = self.aggregate_resolver(resolved_arguments)
         return ResolvedConditionalAggregate(
@@ -354,6 +359,7 @@ def resolve(
         search_type: constants.SearchType,
         resolved_arguments: list[AttributeKey | Any],
         snuba_params: SnubaParams,
+        query_result_cache: dict[str, EAPResponse],
     ) -> ResolvedFormula:
         resolver_settings = ResolverSettings(
             extrapolation_mode=(
@@ -362,6 +368,7 @@ def resolve(
                 else ExtrapolationMode.EXTRAPOLATION_MODE_NONE
             ),
             snuba_params=snuba_params,
+            query_result_cache=query_result_cache,
         )
 
         return ResolvedFormula(

src/sentry/search/eap/resolver.py

@@ -50,7 +50,7 @@
     VirtualColumnDefinition,
 )
 from sentry.search.eap.spans.attributes import SPANS_INTERNAL_TO_PUBLIC_ALIAS_MAPPINGS
-from sentry.search.eap.types import SearchResolverConfig
+from sentry.search.eap.types import EAPResponse, SearchResolverConfig
 from sentry.search.eap.utils import validate_sampling
 from sentry.search.events import constants as qb_constants
 from sentry.search.events import fields
@@ -69,6 +69,7 @@ class SearchResolver:
     config: SearchResolverConfig
     definitions: ColumnDefinitions
     granularity_secs: int | None = None
+    _query_result_cache: dict[str, EAPResponse] = field(default_factory=dict)
     _resolved_attribute_cache: dict[
         str, tuple[ResolvedAttribute, VirtualColumnDefinition | None]
     ] = field(default_factory=dict)
@@ -947,6 +948,7 @@ def resolve_function(self, column: str, match: Match | None = None) -> tuple[
             search_type=search_type,
             resolved_arguments=resolved_arguments,
             snuba_params=self.params,
+            query_result_cache=self._query_result_cache,
         )
 
         resolved_context = None

src/sentry/search/eap/spans/formulas.py

@@ -245,26 +245,37 @@ def failure_rate(_: ResolvedArguments, settings: ResolverSettings) -> Column.Bin
 
 
 def get_count_of_vital(vital: str, settings: ResolverSettings) -> float:
+    cache_key = "totalvitalcounts"
+    response = None
+    vital_column = f"count_scores(measurements.score.{vital})"
+
+    if cache_key in settings["query_result_cache"]:
+        response = settings["query_result_cache"][cache_key]
+
+    else:
+        snuba_params = settings["snuba_params"]
+        query_string = snuba_params.query_string
+
+        vital_columns = [f"count_scores({v})" for v in WEB_VITALS_MEASUREMENTS]
+
+        response = spans_rpc.run_table_query(
+            snuba_params,
+            query_string=query_string if query_string is not None else "",
+            referrer=cache_key,
+            selected_columns=vital_columns,
+            orderby=None,
+            offset=0,
+            limit=1,
+            sampling_mode=snuba_params.sampling_mode,
+            config=SearchResolverConfig(
+                auto_fields=True,
+            ),
+        )
 
-    snuba_params = settings["snuba_params"]
-    query_string = snuba_params.query_string
-
-    rpc_res = spans_rpc.run_table_query(
-        snuba_params,
-        query_string=query_string if query_string is not None else "",
-        referrer=f"totalvitalcount_{vital}",
-        selected_columns=[f"count_scores(measurements.score.{vital}) as count"],
-        orderby=None,
-        offset=0,
-        limit=1,
-        sampling_mode=snuba_params.sampling_mode,
-        config=SearchResolverConfig(
-            auto_fields=True,
-        ),
-    )
+        settings["query_result_cache"][cache_key] = response
 
-    if len(rpc_res["data"]) > 0 and rpc_res["data"][0]["count"] is not None:
-        return rpc_res["data"][0]["count"]
+    if len(response["data"]) > 0 and response["data"][0][vital_column] is not None:
+        return response["data"][0][vital_column]
 
     return 0