feat(metrics): Metrics interface to run queries from code (#33735)

In order to make metrics queries from code, transform QueryDefinition
into a dataclass similar to snuba_sdk.Query, such that queries can be
composed programmatically.

mypy.ini

@@ -77,6 +77,7 @@ files = src/sentry/analytics/,
         src/sentry/snuba/metrics/fields/histogram.py,
         src/sentry/snuba/metrics/fields/base.py,
         src/sentry/snuba/metrics/naming_layer/,
+        src/sentry/snuba/metrics/query.py,
         src/sentry/spans/,
         src/sentry/tasks/app_store_connect.py,
         src/sentry/tasks/low_priority_symbolication.py,

src/sentry/api/endpoints/organization_metrics.py

@@ -8,7 +8,7 @@ from sentry.api.exceptions import ResourceDoesNotExist
 from sentry.api.paginator import GenericOffsetPaginator
 from sentry.api.utils import InvalidParams
 from sentry.snuba.metrics import (
-    QueryDefinition,
+    APIQueryDefinition,
     get_metrics,
     get_series,
     get_single_metric_info,
@@ -117,10 +117,11 @@ class OrganizationMetricsDataEndpoint(OrganizationEndpoint):
 
         def data_fn(offset: int, limit: int):
             try:
-                query = QueryDefinition(
-                    request.GET, paginator_kwargs={"limit": limit, "offset": offset}
+                query = APIQueryDefinition(
+                    projects, request.GET, paginator_kwargs={"limit": limit, "offset": offset}
                 )
-                data = get_series(projects, query)
+                data = get_series(projects, query.to_query_definition())
+                data["query"] = query.query
             except (
                 InvalidField,
                 InvalidParams,

src/sentry/release_health/duplex.py

@@ -619,7 +619,7 @@ def get_sessionsv2_schema(now: datetime, query: QueryDefinition) -> Mapping[str,
                 # timestamp 09:00 contains data for the range 09:00 - 10:00,
                 # And we want to still exclude that at 10:01
                 comparator if timestamp < max_timestamp else ComparatorType.Ignore
-                for timestamp in get_intervals(query)
+                for timestamp in get_intervals(query.start, query.end, query.rollup)
             ]
         )
         for field, comparator in schema_for_totals.items()

src/sentry/release_health/metrics_sessions_v2.py

@@ -711,7 +711,7 @@ def run_sessions_query(
 
     data_points = _flatten_data(org_id, data)
 
-    intervals = list(get_intervals(query_clone))
+    intervals = list(get_intervals(query_clone.start, query_clone.end, query_clone.rollup))
     timestamp_index = {timestamp.isoformat(): index for index, timestamp in enumerate(intervals)}
 
     def default_for(field: SessionsQueryFunction) -> SessionsQueryValue:

src/sentry/snuba/metrics/__init__.py

@@ -1,4 +1,5 @@
 from .datasource import *  # NOQA
 from .fields import *  # NOQA
+from .query import *  # NOQA
 from .query_builder import *  # NOQA
 from .utils import *  # NOQA

src/sentry/snuba/metrics/datasource.py

@@ -10,6 +10,7 @@ __all__ = ("get_metrics", "get_tags", "get_tag_values", "get_series", "get_singl
 import logging
 from collections import defaultdict, deque
 from copy import copy
+from dataclasses import replace
 from operator import itemgetter
 from typing import Any, Dict, Mapping, Optional, Sequence, Set, Tuple, Union
 
@@ -23,9 +24,9 @@ from sentry.snuba.dataset import EntityKey
 from sentry.snuba.metrics.fields import run_metrics_query
 from sentry.snuba.metrics.fields.base import get_derived_metrics, org_id_from_projects
 from sentry.snuba.metrics.naming_layer.mapping import get_mri, get_public_name_from_mri
+from sentry.snuba.metrics.query import QueryDefinition
 from sentry.snuba.metrics.query_builder import (
     ALLOWED_GROUPBY_COLUMNS,
-    QueryDefinition,
     SnubaQueryBuilder,
     SnubaResultConverter,
     get_intervals,
@@ -421,16 +422,15 @@ def get_tag_values(
 
 def get_series(projects: Sequence[Project], query: QueryDefinition) -> dict:
     """Get time series for the given query"""
-    intervals = list(get_intervals(query))
+    intervals = list(get_intervals(query.start, query.end, query.granularity.granularity))
     results = {}
-    org_id = org_id_from_projects(projects)
     fields_in_entities = {}
 
     if not query.groupby:
         # When there is no groupBy columns specified, we don't want to go through running an
         # initial query first to get the groups because there are no groups, and it becomes just
         # one group which is basically identical to eliminating the orderBy altogether
-        query.orderby = None
+        query = replace(query, orderby=None)
 
     if query.orderby is not None:
         # ToDo(ahmed): Now that we have conditional aggregates as select statements, we might be
@@ -445,14 +445,12 @@ def get_series(projects: Sequence[Project], query: QueryDefinition) -> dict:
 
         # Multi-field select with order by functionality. Currently only supports the
         # performance table.
-        original_query_fields = copy(query.fields)
+        original_select = copy(query.select)
 
         # The initial query has to contain only one field which is the same as the order by
         # field
-        orderby_expr, orderby_parsed = [
-            (key, value) for key, value in query.fields.items() if value == query.orderby[0]
-        ][0]
-        query.fields = {orderby_expr: orderby_parsed}
+        orderby_field = [field for field in query.select if field == query.orderby.field][0]
+        query = replace(query, select=[orderby_field])
 
         snuba_queries, _ = SnubaQueryBuilder(projects, query).get_snuba_queries()
         if len(snuba_queries) > 1:
@@ -485,8 +483,7 @@ def get_series(projects: Sequence[Project], query: QueryDefinition) -> dict:
             # the group by tags from the first query so we basically remove the order by columns,
             # and reset the query fields to the original fields because in the second query,
             # we want to query for all the metrics in the request api call
-            query.orderby = None
-            query.fields = original_query_fields
+            query = replace(query, select=original_select, orderby=None)
 
             query_builder = SnubaQueryBuilder(projects, query)
             snuba_queries, fields_in_entities = query_builder.get_snuba_queries()
@@ -495,8 +492,10 @@ def get_series(projects: Sequence[Project], query: QueryDefinition) -> dict:
             # will be used to filter down and order the results of the 2nd query.
             # For example, (project_id, transaction) is translated to (project_id, tags[3])
             groupby_tags = tuple(
-                resolve_tag_key(org_id, field) if field not in ALLOWED_GROUPBY_COLUMNS else field
-                for field in query.groupby
+                resolve_tag_key(query.org_id, field)
+                if field not in ALLOWED_GROUPBY_COLUMNS
+                else field
+                for field in (query.groupby or [])
             )
 
             # Dictionary that contains the conditions that are required to be added to the where
@@ -609,7 +608,6 @@ def get_series(projects: Sequence[Project], query: QueryDefinition) -> dict:
     return {
         "start": query.start,
         "end": query.end,
-        "query": query.query,
         "intervals": intervals,
         "groups": converter.translate_results(),
     }

src/sentry/snuba/metrics/fields/base.py

@@ -60,6 +60,7 @@ from sentry.snuba.metrics.fields.snql import (
 )
 from sentry.snuba.metrics.naming_layer.mapping import get_public_name_from_mri
 from sentry.snuba.metrics.naming_layer.mri import SessionMRI, TransactionMRI
+from sentry.snuba.metrics.query import QueryDefinition
 from sentry.snuba.metrics.utils import (
     DEFAULT_AGGREGATES,
     GRANULARITY,
@@ -94,9 +95,6 @@ __all__ = (
 SnubaDataType = Dict[str, Any]
 PostQueryFuncReturnType = Optional[Union[Tuple[Any, ...], ClickhouseHistogram, int, float]]
 
-if TYPE_CHECKING:
-    from sentry.snuba.metrics.query_builder import QueryDefinition
-
 
 def run_metrics_query(
     *,
@@ -1064,7 +1062,7 @@ DERIVED_METRICS: Mapping[str, DerivedMetricExpression] = {
 }
 
 
-DERIVED_OPS: Mapping[str, DerivedOp] = {
+DERIVED_OPS: Mapping[MetricOperationType, DerivedOp] = {
     derived_op.op: derived_op
     for derived_op in [
         DerivedOp(
@@ -1090,7 +1088,9 @@ DERIVED_ALIASES: Mapping[str, AliasedDerivedMetric] = {
 }
 
 
-def metric_object_factory(op: Optional[str], metric_mri: str) -> MetricExpressionBase:
+def metric_object_factory(
+    op: Optional[MetricOperationType], metric_mri: str
+) -> MetricExpressionBase:
     """Returns an appropriate instance of MetricsFieldBase object"""
     if op in DERIVED_OPS and metric_mri in DERIVED_METRICS:
         raise InvalidParams("derived ops cannot be used on derived metrics")

src/sentry/snuba/metrics/query.py

@@ -0,0 +1,51 @@
+""" Classes needed to build a metrics query. Inspired by snuba_sdk.query. """
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Literal, Optional, Sequence, Union
+
+from snuba_sdk import Direction, Granularity, Limit, Offset
+from snuba_sdk.conditions import ConditionGroup
+
+from .utils import MetricOperationType
+
+# TODO: Add __all__ to be consistent with sibling modules
+
+
+@dataclass(frozen=True)
+class MetricField:
+    op: Optional[MetricOperationType]
+    metric_name: str
+
+
+Tag = str
+Groupable = Union[Tag, Literal["project_id"]]
+
+
+@dataclass(frozen=True)
+class OrderBy:
+    field: MetricField
+    direction: Direction
+
+
+@dataclass(frozen=True)
+class QueryDefinition:
+    """Definition of a metrics query, inspired by snuba_sdk.Query"""
+
+    org_id: int
+    project_ids: Sequence[int]
+    select: Sequence[MetricField]
+    start: datetime
+    end: datetime
+    granularity: Granularity
+    where: Optional[ConditionGroup] = None  # TODO: Should restrict
+    groupby: Optional[Sequence[Groupable]] = None
+    orderby: Optional[OrderBy] = None
+    limit: Optional[Limit] = None
+    offset: Optional[Offset] = None
+    include_totals: bool = True
+    include_series: bool = True
+
+    # TODO: These should be properties of the Histogram field
+    histogram_buckets: int = 100
+    histogram_from: Optional[float] = None
+    histogram_to: Optional[float] = None

src/sentry/snuba/metrics/query_builder.py

@@ -1,5 +1,5 @@
 __all__ = (
-    "QueryDefinition",
+    "APIQueryDefinition",
     "SnubaQueryBuilder",
     "SnubaResultConverter",
     "get_date_range",
@@ -35,6 +35,9 @@ from sentry.snuba.metrics.fields.base import (
     org_id_from_projects,
 )
 from sentry.snuba.metrics.naming_layer.mapping import get_mri, get_public_name_from_mri
+from sentry.snuba.metrics.query import MetricField
+from sentry.snuba.metrics.query import OrderBy as MetricsOrderBy
+from sentry.snuba.metrics.query import QueryDefinition, Tag
 from sentry.snuba.metrics.utils import (
     ALLOWED_GROUPBY_COLUMNS,
     FIELD_REGEX,
@@ -46,7 +49,6 @@ from sentry.snuba.metrics.utils import (
     UNALLOWED_TAGS,
     DerivedMetricParseException,
     MetricDoesNotExistException,
-    TimeRange,
 )
 from sentry.snuba.sessions_v2 import ONE_DAY  # TODO: unite metrics and sessions_v2
 from sentry.snuba.sessions_v2 import AllowedResolution, InvalidField, finite_or_none
@@ -54,14 +56,15 @@ from sentry.utils.dates import parse_stats_period, to_datetime, to_timestamp
 from sentry.utils.snuba import parse_snuba_datetime
 
 
-def parse_field(field: str) -> Tuple[Optional[str], str]:
+def parse_field(field: str, query_params) -> MetricField:
     derived_metrics_mri = get_derived_metrics(exclude_private=True)
     matches = FIELD_REGEX.match(field)
     try:
         if matches is None:
             raise TypeError
         operation = matches[1]
-        metric_mri = get_mri(matches[2])
+        metric_name = matches[2]
+        metric_mri = get_mri(metric_name)
         if metric_mri in derived_metrics_mri and isinstance(
             derived_metrics_mri[metric_mri], DerivedMetricExpression
         ):
@@ -70,12 +73,13 @@ def parse_field(field: str) -> Tuple[Optional[str], str]:
                 f"already a derived metric with an aggregation applied to it."
             )
     except (IndexError, TypeError):
-        metric_mri = get_mri(field)
+        metric_name = field
+        metric_mri = get_mri(metric_name)
         if metric_mri in derived_metrics_mri and isinstance(
             derived_metrics_mri[metric_mri], DerivedMetricExpression
         ):
             # The isinstance check is there to foreshadow adding raw metric aliases
-            return None, metric_mri
+            return MetricField(op=None, metric_name=metric_name)
         raise InvalidField(
             f"Failed to parse '{field}'. Must be something like 'sum(my_metric)', or a supported "
             f"aggregate derived metric like `session.crash_free_rate"
@@ -86,7 +90,7 @@ def parse_field(field: str) -> Tuple[Optional[str], str]:
             f"Invalid operation '{operation}'. Must be one of {', '.join(OPERATIONS)}"
         )
 
-    return operation, metric_mri
+    return MetricField(operation, metric_name)
 
 
 def resolve_tags(org_id: int, input_: Any) -> Any:
@@ -149,7 +153,7 @@ def parse_query(query_string: str) -> Sequence[Condition]:
     return where
 
 
-class QueryDefinition:
+class APIQueryDefinition:
     """
     This is the definition of the query the user wants to execute.
     This is constructed out of the request params, and also contains a list of
@@ -159,7 +163,8 @@ class QueryDefinition:
 
     """
 
-    def __init__(self, query_params, paginator_kwargs: Optional[Dict] = None):
+    def __init__(self, projects, query_params, paginator_kwargs: Optional[Dict] = None):
+        self._projects = projects
         paginator_kwargs = paginator_kwargs or {}
 
         self.query = query_params.get("query", "")
@@ -170,14 +175,12 @@ class QueryDefinition:
         if len(raw_fields) == 0:
             raise InvalidField('Request is missing a "field"')
 
-        self.fields = {}
+        self.fields = []
         for key in raw_fields:
-            op, metric_mri = parse_field(key)
-            mri_key = f"{op}({metric_mri})" if op is not None else metric_mri
-            self.fields[mri_key] = (op, metric_mri)
+            self.fields.append(parse_field(key, query_params))
 
         self.orderby = self._parse_orderby(query_params)
-        self.limit = self._parse_limit(query_params, paginator_kwargs)
+        self.limit: Optional[Limit] = self._parse_limit(query_params, paginator_kwargs)
         self.offset = self._parse_offset(query_params, paginator_kwargs)
 
         start, end, rollup = get_date_range(query_params)
@@ -196,9 +199,32 @@ class QueryDefinition:
         self.include_series = query_params.get("includeSeries", "1") == "1"
         self.include_totals = query_params.get("includeTotals", "1") == "1"
 
+        if not (self.include_series or self.include_totals):
+            raise InvalidParams("Cannot omit both series and totals")
+
         # Validates that time series limit will not exceed the snuba limit of 10,000
         self._validate_series_limit(query_params)
 
+    def to_query_definition(self) -> QueryDefinition:
+        return QueryDefinition(
+            org_id=org_id_from_projects(self._projects),
+            project_ids=[project.id for project in self._projects],
+            include_totals=self.include_totals,
+            include_series=self.include_series,
+            select=self.fields,
+            start=self.start,
+            end=self.end,
+            where=self.parsed_query,
+            groupby=self.groupby,
+            orderby=self.orderby,
+            limit=self.limit,
+            offset=self.offset,
+            granularity=Granularity(self.rollup),
+            histogram_buckets=self.histogram_buckets,
+            histogram_from=self.histogram_from,
+            histogram_to=self.histogram_to,
+        )
+
     def _parse_orderby(self, query_params):
         orderby = query_params.getlist("orderBy", [])
         if not orderby:
@@ -213,16 +239,16 @@ class QueryDefinition:
             direction = Direction.DESC
 
         try:
-            op, metric_mri = parse_field(orderby)
+            field = parse_field(orderby, query_params)
         except KeyError:
             # orderBy one of the group by fields may be supported in the future
             raise InvalidParams("'orderBy' must be one of the provided 'fields'")
 
-        return (op, metric_mri), direction
+        return MetricsOrderBy(field, direction)
 
     def _parse_limit(self, query_params, paginator_kwargs):
         if self.orderby:
-            return paginator_kwargs.get("limit")
+            return Limit(paginator_kwargs.get("limit"))
         else:
             per_page = query_params.get("per_page")
             if per_page is not None:
@@ -246,19 +272,20 @@ class QueryDefinition:
 
     def _validate_series_limit(self, query_params):
         if self.limit:
-            if (self.end - self.start).total_seconds() / self.rollup * self.limit > MAX_POINTS:
+            if (
+                self.end - self.start
+            ).total_seconds() / self.rollup * self.limit.limit > MAX_POINTS:
                 raise InvalidParams(
                     f"Requested interval of {query_params.get('interval', '1h')} with statsPeriod of "
                     f"{query_params.get('statsPeriod')} is too granular for a per_page of "
-                    f"{self.limit} elements. Increase your interval, decrease your statsPeriod, "
+                    f"{self.limit.limit} elements. Increase your interval, decrease your statsPeriod, "
                     f"or decrease your per_page parameter."
                 )
 
 
-def get_intervals(query: TimeRange):
-    start = query.start
-    end = query.end
-    delta = timedelta(seconds=query.rollup)
+def get_intervals(start: datetime, end: datetime, granularity: int):
+    assert granularity > 0
+    delta = timedelta(seconds=granularity)
     while start < end:
         yield start
         start += delta
@@ -324,18 +351,17 @@ class SnubaQueryBuilder:
 
     def __init__(self, projects: Sequence[Project], query_definition: QueryDefinition):
         self._projects = projects
-        self._org_id = org_id_from_projects(projects)
         self._query_definition = query_definition
+        self._org_id = query_definition.org_id
 
     def _build_where(self) -> List[Union[BooleanCondition, Condition]]:
-        assert self._projects
         where: List[Union[BooleanCondition, Condition]] = [
             Condition(Column("org_id"), Op.EQ, self._org_id),
-            Condition(Column("project_id"), Op.IN, [p.id for p in self._projects]),
+            Condition(Column("project_id"), Op.IN, self._query_definition.project_ids),
             Condition(Column(TS_COL_QUERY), Op.GTE, self._query_definition.start),
             Condition(Column(TS_COL_QUERY), Op.LT, self._query_definition.end),
         ]
-        filter_ = resolve_tags(self._org_id, self._query_definition.parsed_query)
+        filter_ = resolve_tags(self._org_id, self._query_definition.where)
         if filter_:
             where.extend(filter_)
 
@@ -343,30 +369,32 @@ class SnubaQueryBuilder:
 
     def _build_groupby(self) -> List[Column]:
         groupby_cols = []
-        for field in self._query_definition.groupby:
+        for field in self._query_definition.groupby or []:
             if field in UNALLOWED_TAGS:
                 raise InvalidParams(f"Tag name {field} cannot be used to groupBy query")
             if field in ALLOWED_GROUPBY_COLUMNS:
                 groupby_cols.append(Column(field))
             else:
+                assert isinstance(field, Tag)
                 groupby_cols.append(Column(resolve_tag_key(self._org_id, field)))
         return groupby_cols
 
     def _build_orderby(self) -> Optional[List[OrderBy]]:
         if self._query_definition.orderby is None:
             return None
-        (op, metric_mri), direction = self._query_definition.orderby
+        orderby = self._query_definition.orderby
+        op = orderby.field.op
+        metric_mri = get_mri(orderby.field.metric_name)
         metric_field_obj = metric_object_factory(op, metric_mri)
         return metric_field_obj.generate_orderby_clause(
-            projects=self._projects, direction=direction, query_definition=self._query_definition
+            projects=self._projects,
+            direction=orderby.direction,
+            query_definition=self._query_definition,
         )
 
     def __build_totals_and_series_queries(
         self, entity, select, where, groupby, orderby, limit, offset, rollup, intervals_len
     ):
-        if not self._query_definition.include_totals and not self._query_definition.include_series:
-            return {}
-
         rv = {}
         totals_query = Query(
             dataset=Dataset.Metrics.value,
@@ -374,9 +402,9 @@ class SnubaQueryBuilder:
             groupby=groupby,
             select=select,
             where=where,
-            limit=Limit(limit or MAX_POINTS),
+            limit=limit or Limit(MAX_POINTS),
             offset=Offset(offset or 0),
-            granularity=Granularity(rollup),
+            granularity=rollup,
             orderby=orderby,
         )
 
@@ -391,7 +419,7 @@ class SnubaQueryBuilder:
             # In a series query, we also need to factor in the len of the intervals array
             series_limit = MAX_POINTS
             if limit:
-                series_limit = limit * intervals_len
+                series_limit = limit.limit * intervals_len
             rv["series"] = series_query.set_limit(series_limit)
 
         return rv
@@ -416,8 +444,9 @@ class SnubaQueryBuilder:
         metric_mri_to_obj_dict = {}
         fields_in_entities = {}
 
-        for op, metric_mri in self._query_definition.fields.values():
-            metric_field_obj = metric_object_factory(op, metric_mri)
+        for field in self._query_definition.select:
+            metric_mri = get_mri(field.metric_name)
+            metric_field_obj = metric_object_factory(field.op, metric_mri)
             # `get_entity` is called the first, to fetch the entities of constituent metrics,
             # and validate especially in the case of SingularEntityDerivedMetric that it is
             # actually composed of metrics that belong to the same entity
@@ -455,8 +484,8 @@ class SnubaQueryBuilder:
             if entity not in self._implemented_datasets:
                 raise NotImplementedError(f"Dataset not yet implemented: {entity}")
 
-            metric_mri_to_obj_dict[(op, metric_mri)] = metric_field_obj
-            fields_in_entities.setdefault(entity, []).append((op, metric_mri))
+            metric_mri_to_obj_dict[(field.op, metric_mri)] = metric_field_obj
+            fields_in_entities.setdefault(entity, []).append((field.op, metric_mri))
 
         where = self._build_where()
         groupby = self._build_groupby()
@@ -465,8 +494,8 @@ class SnubaQueryBuilder:
         for entity, fields in fields_in_entities.items():
             select = []
             metric_ids_set = set()
-            for op, name in fields:
-                metric_field_obj = metric_mri_to_obj_dict[(op, name)]
+            for field in fields:
+                metric_field_obj = metric_mri_to_obj_dict[field]
                 select += metric_field_obj.generate_select_statements(
                     projects=self._projects, query_definition=self._query_definition
                 )
@@ -489,8 +518,16 @@ class SnubaQueryBuilder:
                 orderby=orderby,
                 limit=self._query_definition.limit,
                 offset=self._query_definition.offset,
-                rollup=self._query_definition.rollup,
-                intervals_len=len(list(get_intervals(self._query_definition))),
+                rollup=self._query_definition.granularity,
+                intervals_len=len(
+                    list(
+                        get_intervals(
+                            self._query_definition.start,
+                            self._query_definition.end,
+                            self._query_definition.granularity.granularity,
+                        )
+                    )
+                ),
             )
 
         return queries_dict, fields_in_entities
@@ -513,7 +550,9 @@ class SnubaResultConverter:
         self._query_definition = query_definition
 
         # This is a set of all the `(op, metric_mri)` combinations passed in the query_definition
-        self._query_definition_fields_set = set(query_definition.fields.values())
+        self._query_definition_fields_set = {
+            (field.op, get_mri(field.metric_name)) for field in query_definition.select
+        }
         # This is a set of all queryable `(op, metric_mri)` combinations. Queryable can mean it
         # includes one of the following: AggregatedRawMetric (op, metric_mri), instance of
         # SingularEntityDerivedMetric or the instances of SingularEntityDerivedMetric that are

src/sentry/snuba/metrics/utils.py

@@ -26,7 +26,6 @@ __all__ = (
     "MetricDoesNotExistException",
     "MetricDoesNotExistInIndexer",
     "NotSupportedOverCompositeEntityException",
-    "TimeRange",
     "MetricEntity",
     "UNALLOWED_TAGS",
     "combine_dictionary_of_list_values",
@@ -35,7 +34,6 @@ __all__ = (
 
 import re
 from abc import ABC
-from datetime import datetime
 from typing import (
     Collection,
     Dict,
@@ -43,7 +41,6 @@ from typing import (
     Literal,
     Mapping,
     Optional,
-    Protocol,
     Sequence,
     Tuple,
     TypedDict,
@@ -64,7 +61,18 @@ TAG_REGEX = re.compile(r"^(\w|\.|_)+$")
 
 #: A function that can be applied to a metric
 MetricOperationType = Literal[
-    "avg", "count", "max", "min", "p50", "p75", "p90", "p95", "p99", "histogram"
+    "avg",
+    "count",
+    "count_unique",
+    "sum",
+    "max",
+    "min",
+    "p50",
+    "p75",
+    "p90",
+    "p95",
+    "p99",
+    "histogram",
 ]
 MetricUnit = Literal["seconds"]
 #: The type of metric, which determines the snuba entity to query
@@ -72,15 +80,14 @@ MetricType = Literal["counter", "set", "distribution", "numeric"]
 
 MetricEntity = Literal["metrics_counters", "metrics_sets", "metrics_distributions"]
 
-OP_TO_SNUBA_FUNCTION = {
+OP_TO_SNUBA_FUNCTION: Mapping[str, Mapping[MetricOperationType, str]] = {
     "metrics_counters": {"sum": "sumIf"},
     "metrics_distributions": {
         "avg": "avgIf",
         "count": "countIf",
         "max": "maxIf",
         "min": "minIf",
-        # TODO: Would be nice to use `quantile(0.50)` (singular) here, but snuba responds with an error
-        "p50": "quantilesIf(0.50)",
+        "p50": "quantilesIf(0.50)",  # TODO: Would be nice to use `quantile(0.50)` (singular) here, but snuba responds with an error
         "p75": "quantilesIf(0.75)",
         "p90": "quantilesIf(0.90)",
         "p95": "quantilesIf(0.95)",
@@ -199,9 +206,3 @@ class DerivedMetricParseException(DerivedMetricException):
 
 class NotSupportedOverCompositeEntityException(DerivedMetricException):
     ...
-
-
-class TimeRange(Protocol):
-    start: datetime
-    end: datetime
-    rollup: int