typing(metrics): Various fixes (#60614)

pyproject.toml

@@ -564,8 +564,6 @@ module = [
     "sentry.snuba.metrics.datasource",
     "sentry.snuba.metrics.fields.base",
     "sentry.snuba.metrics.query_builder",
-    "sentry.snuba.metrics_enhanced_performance",
-    "sentry.snuba.metrics_performance",
     "sentry.snuba.sessions",
     "sentry.snuba.sessions_v2",
     "sentry.snuba.spans_indexed",

src/sentry/search/events/types.py

@@ -9,7 +9,7 @@ from snuba_sdk.conditions import BooleanCondition, Condition
 from snuba_sdk.entity import Entity
 from snuba_sdk.function import CurriedFunction, Function
 from snuba_sdk.orderby import OrderBy
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from sentry.models.environment import Environment
 from sentry.models.organization import Organization
@@ -41,6 +41,7 @@ class QueryFramework:
 class EventsMeta(TypedDict):
     fields: Dict[str, str]
     tips: Dict[str, str]
+    isMetricsData: NotRequired[bool]
 
 
 class EventsResponse(TypedDict):

src/sentry/snuba/discover.py

@@ -10,7 +10,7 @@ import sentry_sdk
 from sentry_relay.consts import SPAN_STATUS_CODE_TO_NAME
 from snuba_sdk.conditions import Condition, Op
 from snuba_sdk.function import Function
-from typing_extensions import TypedDict
+from typing_extensions import NotRequired, TypedDict
 
 from sentry.discover.arithmetic import categorize_columns
 from sentry.exceptions import InvalidSearchQuery
@@ -69,6 +69,17 @@ FacetResult = namedtuple("FacetResult", ["key", "value", "count"])
 
 class EventsMeta(TypedDict):
     fields: Dict[str, str]
+    datasetReason: NotRequired[str]
+    isMetricsData: NotRequired[bool]
+    isMetricsExtractedData: NotRequired[bool]
+
+
+# When calling make build-spectacular-docs we hit this issue
+# https://github.com/tfranzel/drf-spectacular/issues/1041
+# This is a work around
+EventsMeta.__annotations__["datasetReason"] = str
+EventsMeta.__annotations__["isMetricsData"] = bool
+EventsMeta.__annotations__["isMetricsExtractedData"] = bool
 
 
 class EventsResponse(TypedDict):

src/sentry/snuba/metrics_performance.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from datetime import timedelta
 from typing import Any, Dict, List, Optional, Sequence
@@ -14,7 +16,7 @@ from sentry.search.events.builder import (
     TopMetricsQueryBuilder,
 )
 from sentry.search.events.fields import get_function_alias
-from sentry.search.events.types import QueryBuilderConfig
+from sentry.search.events.types import EventsResponse, QueryBuilderConfig
 from sentry.snuba import discover
 from sentry.snuba.dataset import Dataset
 from sentry.snuba.metrics.extraction import MetricSpecType
@@ -100,7 +102,7 @@ def bulk_timeseries_query(
     on_demand_metrics_type: Optional[MetricSpecType] = None,
     groupby: Optional[Column] = None,
     apply_formatting: Optional[bool] = True,
-) -> SnubaTSResult:
+) -> SnubaTSResult | EventsResponse:
     """
     High-level API for doing *bulk* arbitrary user timeseries queries against events.
     this API should match that of sentry.snuba.discover.timeseries_query
@@ -113,7 +115,6 @@ def bulk_timeseries_query(
     if metrics_compatible:
         with sentry_sdk.start_span(op="mep", description="TimeseriesMetricQueryBuilder"):
             metrics_queries = []
-            metrics_query = None
             for query in queries:
                 metrics_query = TimeseriesMetricQueryBuilder(
                     params,
@@ -133,18 +134,18 @@ def bulk_timeseries_query(
 
             metrics_referrer = referrer + ".metrics-enhanced"
             bulk_result = bulk_snql_query(metrics_queries, metrics_referrer)
-            result = {"data": []}
+            _result: dict[str, Any] = {"data": []}
             for br in bulk_result:
-                result["data"] = [*result["data"], *br["data"]]
-                result["meta"] = br["meta"]
+                _result["data"] = [*_result["data"], *br["data"]]
+                _result["meta"] = br["meta"]
         with sentry_sdk.start_span(op="mep", description="query.transform_results"):
-            result = metrics_query.process_results(result)
+            result = metrics_query.process_results(_result)
             sentry_sdk.set_tag("performance.dataset", "metrics")
             result["meta"]["isMetricsData"] = True
 
             # Sometimes additional formatting needs to be done downstream
             if not apply_formatting:
-                return result
+                return result  # EventsResponseData type
 
             result["data"] = (
                 discover.zerofill(
@@ -416,7 +417,7 @@ def top_events_timeseries(
 
     translated_groupby = top_events_builder.translated_groupby
 
-    results = (
+    results: dict[str, Any] = (
         {discover.OTHER_KEY: {"order": limit, "data": other_result["data"]}}
         if len(other_result.get("data", []))
         else {}
@@ -560,7 +561,7 @@ def normalize_histogram_results(fields, histogram_params, results):
     """
 
     # zerofill and rename the columns while making sure to adjust for precision
-    bucket_maps = {field: {} for field in fields}
+    bucket_maps: dict[str, Any] = {field: {} for field in fields}
     # Only one row in metrics result
     data = results["data"][0]
     for field in fields:
@@ -568,7 +569,7 @@ def normalize_histogram_results(fields, histogram_params, results):
         histogram_alias = get_function_alias(histogram_column)
         bucket_maps[field] = {start: height for start, end, height in data[histogram_alias]}
 
-    new_data = {field: [] for field in fields}
+    new_data: dict[str, Any] = {field: [] for field in fields}
     for i in range(histogram_params.num_buckets):
         bucket = histogram_params.start_offset + histogram_params.bucket_size * i
         for field in fields:

src/sentry/utils/snuba.py

@@ -810,7 +810,7 @@ def bulk_snql_query(
     requests: List[Request],
     referrer: Optional[str] = None,
     use_cache: bool = False,
-) -> Mapping[str, Any]:
+) -> ResultSet:
     # XXX (evanh): This function does none of the extra processing that the
     # other functions do here. It does not add any automatic conditions, format
     # results, nothing. Use at your own risk.
@@ -1050,7 +1050,6 @@ def query(
     use_cache=False,
     **kwargs,
 ):
-
     aggregations = aggregations or [["count()", "", "aggregate"]]
     filter_keys = filter_keys or {}
     selected_columns = selected_columns or []
@@ -1190,7 +1189,7 @@ def resolve_condition(cond, column_resolver):
             return cond
 
         func_args = cond[index + 1]
-        for (i, arg) in enumerate(func_args):
+        for i, arg in enumerate(func_args):
             # Nested function
             try:
                 if isinstance(arg, (list, tuple)):
@@ -1274,7 +1273,7 @@ def aliased_query_params(
     derived_columns = []
     resolve_func = resolve_column(dataset)
     if selected_columns:
-        for (i, col) in enumerate(selected_columns):
+        for i, col in enumerate(selected_columns):
             if isinstance(col, (list, tuple)):
                 derived_columns.append(col[2])
             else:
@@ -1298,7 +1297,7 @@ def aliased_query_params(
     if orderby:
         # Don't mutate in case we have a default order passed.
         updated_order = []
-        for (i, order) in enumerate(orderby):
+        for i, order in enumerate(orderby):
             order_field = order.lstrip("-")
             if order_field not in derived_columns:
                 order_field = resolve_func(order_field)