Skip to content

[DRAFT] typing: strongly type sentry.snuba.metrics.query_builder #121039

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cursor wants to merge 2 commits into
base: master
Choose a base branch
from
+91 −62
Draft

[DRAFT] typing: strongly type sentry.snuba.metrics.query_builder #121039

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c86f98c
typing: strongly type metrics query_builder module
cursoragent Aug 2, 2026
18cf7cb
typing: preserve metric operation params behavior
cursoragent Aug 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -388,7 +388,6 @@ ignore_missing_imports = true
# - python3 -m tools.mypy_helpers.find_easiest_modules
[[tool.mypy.overrides]]
module = [
"sentry.snuba.metrics.query_builder",
"sentry.testutils.cases",
]
disable_error_code = [
Expand Down
152 changes: 91 additions & 61 deletions src/sentry/snuba/metrics/query_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, TypedDict, overload
from typing import Any, TypedDict, cast, overload

import sentry_sdk
from snuba_sdk import (
Expand Down Expand Up @@ -42,6 +42,7 @@
from sentry.snuba.metrics.fields import metric_object_factory
from sentry.snuba.metrics.fields.base import (
COMPOSITE_ENTITY_CONSTITUENT_ALIAS,
MetricOperationParams,
MetricExpressionBase,
generate_bottom_up_dependency_tree_for_metrics,
org_id_from_projects,
Expand Down Expand Up @@ -71,6 +72,8 @@
DerivedMetricParseException,
MetricDoesNotExistException,
MetricEntity,
MetricOperationType,
OPERATIONS,
get_num_intervals,
get_timestamp_column_name,
require_rhs_condition_resolution,
Expand All @@ -93,6 +96,15 @@


QUERY_PROJECT_LIMIT = 10
MetricFieldKey = tuple[MetricOperationType | None, str, str]


def _normalize_metric_operation_params(
params: dict[str, None | str | int | float | Sequence[tuple[str | int, ...]]] | None,
) -> MetricOperationParams | None:
# MetricField params currently support richer shapes than MetricOperationParams
# models. Keep runtime behavior unchanged while preserving type safety at callsites.
return cast(MetricOperationParams | None, params)


def _strip_project_id(condition: Condition) -> Condition | None:
Expand Down Expand Up @@ -134,7 +146,10 @@ def parse_public_field(field: str) -> MetricField:
operation = None
metric_name = field

return MetricField(operation, get_mri(metric_name))
if operation is not None and operation not in OPERATIONS:
raise InvalidParams(f"Invalid operation '{operation}'")

return MetricField(cast(MetricOperationType | None, operation), get_mri(metric_name))


def transform_null_transaction_to_unparameterized(use_case_id, org_id, alias=None):
Expand Down Expand Up @@ -747,11 +762,14 @@ def translate_meta_results(
continue
elif alias_type == AliasMetaType.GROUP_BY_METRIC_FIELD:
metric_groupby_field = alias_to_metric_group_by_field[record["name"]]
assert isinstance(metric_groupby_field.field, MetricField)
defined_parent_meta_type = get_metric_object_from_metric_field(
metric_groupby_field.field
).get_meta_type()

record["type"] = defined_parent_meta_type
record["type"] = (
record["type"] if defined_parent_meta_type is None else defined_parent_meta_type
)
elif alias_type == AliasMetaType.TAG:
record["type"] = "string"
elif alias_type == AliasMetaType.DATASET_COLUMN or alias_type == AliasMetaType.TIME_COLUMN:
Expand Down Expand Up @@ -823,82 +841,77 @@ def generate_snql_for_action_by_fields(
function has been designed to share as much logic as possible, however, it should be refactored in case
the snql generation starts to diverge significantly.
"""

is_group_by = isinstance(metric_action_by_field, MetricGroupByField)
is_order_by = isinstance(metric_action_by_field, MetricOrderByField)
if not is_group_by and not is_order_by:
groupby_field: MetricGroupByField | None = None
orderby_field: MetricOrderByField | None = None
if isinstance(metric_action_by_field, MetricGroupByField):
groupby_field = metric_action_by_field
elif isinstance(metric_action_by_field, MetricOrderByField):
orderby_field = metric_action_by_field
else:
raise InvalidParams("The metric action must either be an order by or group by.")

if isinstance(metric_action_by_field.field, str):
action_field = metric_action_by_field.field
if isinstance(action_field, str):
# This transformation is currently supported only for group by because OrderBy doesn't support the Function type.
if is_group_by and metric_action_by_field.field == "transaction":
return transform_null_transaction_to_unparameterized(
use_case_id, org_id, metric_action_by_field.alias
)
if groupby_field is not None and action_field == "transaction":
return transform_null_transaction_to_unparameterized(use_case_id, org_id, groupby_field.alias)

# Handles the case when we are trying to group or order by `project` for example, but we want
# to translate it to `project_id` as that is what the metrics dataset understands.
if metric_action_by_field.field in FIELD_ALIAS_MAPPINGS:
column_name = FIELD_ALIAS_MAPPINGS[metric_action_by_field.field]
elif metric_action_by_field.field in FIELD_ALIAS_MAPPINGS.values():
column_name = metric_action_by_field.field
if action_field in FIELD_ALIAS_MAPPINGS:
column_name = FIELD_ALIAS_MAPPINGS[action_field]
elif action_field in FIELD_ALIAS_MAPPINGS.values():
column_name = action_field
else:
# The support for tags in the order by is disabled for now because there is no need to have it. If the
# need arise, we will implement it.
if is_group_by:
assert isinstance(metric_action_by_field.field, str)
column_name = resolve_tag_key(use_case_id, org_id, metric_action_by_field.field)
if groupby_field is not None:
column_name = resolve_tag_key(use_case_id, org_id, action_field)
else:
raise NotImplementedError(
f"Unsupported string field: {metric_action_by_field.field}"
)
raise NotImplementedError(f"Unsupported string field: {action_field}")

exp = (
AliasedExpression(
exp=Column(name=column_name),
alias=metric_action_by_field.alias,
alias=groupby_field.alias,
)
if is_group_by and not is_column
if groupby_field is not None and not is_column
else Column(name=column_name)
)

if is_order_by:
if orderby_field is not None:
# We return a list in order to use the "extend" method and reduce the number of changes across
# the codebase.
exp = [OrderBy(exp=exp, direction=metric_action_by_field.direction)]
exp = [OrderBy(exp=exp, direction=orderby_field.direction)]

return exp
elif isinstance(metric_action_by_field.field, MetricField):
elif isinstance(action_field, MetricField):
try:
metric_expression = metric_object_factory(
metric_action_by_field.field.op, metric_action_by_field.field.metric_mri
)
metric_expression = metric_object_factory(action_field.op, action_field.metric_mri)

if is_group_by:
if groupby_field is not None:
return metric_expression.generate_groupby_statements(
use_case_id=use_case_id,
alias=metric_action_by_field.field.alias,
params=metric_action_by_field.field.params,
alias=action_field.alias,
params=_normalize_metric_operation_params(action_field.params),
projects=projects,
)[0]
elif is_order_by:
elif orderby_field is not None:
return metric_expression.generate_orderby_clause(
use_case_id=use_case_id,
alias=metric_action_by_field.field.alias,
params=metric_action_by_field.field.params,
alias=action_field.alias,
params=_normalize_metric_operation_params(action_field.params),
projects=projects,
direction=metric_action_by_field.direction,
direction=orderby_field.direction,
)
else:
raise NotImplementedError(
f"Unsupported metric field: {metric_action_by_field.field}"
)
raise NotImplementedError(f"Unsupported metric field: {action_field}")

except IndexError:
raise InvalidParams(f"Cannot resolve {metric_action_by_field.field} into SnQL")
raise InvalidParams(f"Cannot resolve {action_field} into SnQL")
else:
raise NotImplementedError(
f"Unsupported {'group by' if is_group_by else 'order by' if is_order_by else 'None'} field: {metric_action_by_field.field} needs to be either a MetricField or a string"
f"Unsupported {'group by' if groupby_field is not None else 'order by'} field: {action_field} needs to be either a MetricField or a string"
)

def _build_where(self) -> list[BooleanCondition | Condition]:
Expand All @@ -920,21 +933,29 @@ def _build_where(self) -> list[BooleanCondition | Condition]:
metric_expression = metric_object_factory(
condition.lhs.op, condition.lhs.metric_mri
)
lhs_op = condition.lhs.op
resolve_rhs = lhs_op is not None and require_rhs_condition_resolution(lhs_op)
if resolve_rhs and not isinstance(condition.rhs, str):
raise InvalidParams(
"Metric condition rhs should be a string when rhs condition resolution is required."
)
resolved_rhs: int | float | str
if resolve_rhs:
assert isinstance(condition.rhs, str)
resolved_rhs = resolve_tag_value(self._use_case_id, self._org_id, condition.rhs)
else:
resolved_rhs = condition.rhs
try:
metric_condition_filters.append(
Condition(
lhs=metric_expression.generate_where_statements(
use_case_id=self._use_case_id,
params=condition.lhs.params,
params=_normalize_metric_operation_params(condition.lhs.params),
projects=self._projects,
alias=condition.lhs.alias,
)[0],
op=condition.op,
rhs=(
resolve_tag_value(self._use_case_id, self._org_id, condition.rhs)
if require_rhs_condition_resolution(condition.lhs.op)
else condition.rhs
),
rhs=resolved_rhs,
)
)
except IndexError:
Expand Down Expand Up @@ -1069,6 +1090,8 @@ def __build_totals_and_series_queries(
series_limit = self._metrics_query.max_limit

if self._use_case_id in [UseCaseID.TRANSACTIONS, UseCaseID.SPANS]:
if self._metrics_query.interval is None:
raise InvalidParams("Missing interval for discover metrics query.")
time_groupby_column = self.__generate_time_groupby_column_for_discover_queries(
self._metrics_query.interval
)
Expand Down Expand Up @@ -1099,11 +1122,11 @@ def __generate_time_groupby_column_for_discover_queries(interval: int) -> Functi

def __update_query_dicts_with_component_entities(
self,
component_entities: dict[MetricEntity, Sequence[str]],
metric_mri_to_obj_dict: dict[tuple[str | None, str, str], MetricExpressionBase],
fields_in_entities: dict[MetricEntity, list[tuple[str | None, str, str]]],
component_entities: Mapping[MetricEntity | None, Sequence[str]],
metric_mri_to_obj_dict: dict[MetricFieldKey, MetricExpressionBase],
fields_in_entities: dict[MetricEntity, list[MetricFieldKey]],
parent_alias,
) -> dict[tuple[str | None, str, str], MetricExpressionBase]:
) -> dict[MetricFieldKey, MetricExpressionBase]:
# At this point in time, we are only supporting raw metrics in the metrics attribute of
# any instance of DerivedMetric, and so in this case the op will always be None
# ToDo(ahmed): In future PR, we might want to allow for dependency metrics to also have an
Expand All @@ -1117,19 +1140,23 @@ def __update_query_dicts_with_component_entities(
# params. This means that if parent composite metric alias is for example sessions_errored, and it has
# a constituent `e:sessions/error.unique@none` then that constituent will be aliased as
# `e:sessions/error.unique@none__CHILD_OF__sessions_errored`
metric_key = (
metric_key: MetricFieldKey = (
op,
metric_mri,
f"{metric_mri}{COMPOSITE_ENTITY_CONSTITUENT_ALIAS}{parent_alias}",
)
if metric_key not in metric_mri_to_obj_dict:
metric_mri_to_obj_dict[metric_key] = metric_object_factory(op, metric_mri)
if entity is None:
raise DerivedMetricParseException(
"Derived metric component entity cannot be None."
)
fields_in_entities.setdefault(entity, []).append(metric_key)
return metric_mri_to_obj_dict

def get_snuba_queries(self):
metric_mri_to_obj_dict: dict[tuple[str | None, str, str], MetricExpressionBase] = {}
fields_in_entities: dict[MetricEntity, list[tuple[str | None, str, str]]] = {}
metric_mri_to_obj_dict: dict[MetricFieldKey, MetricExpressionBase] = {}
fields_in_entities: dict[MetricEntity, list[MetricFieldKey]] = {}

for select_field in self._metrics_query.select:
metric_field_obj = metric_object_factory(select_field.op, select_field.metric_mri)
Expand Down Expand Up @@ -1201,7 +1228,7 @@ def get_snuba_queries(self):
projects=self._projects,
use_case_id=self._use_case_id,
alias=field[2],
params=params,
params=_normalize_metric_operation_params(params),
)
metric_ids_set |= metric_field_obj.generate_metric_ids(
self._projects, self._use_case_id
Expand Down Expand Up @@ -1255,7 +1282,7 @@ def __init__(
self,
organization_id: int,
metrics_query: DeprecatingMetricsQuery,
fields_in_entities: dict[MetricEntity, list[tuple[str | None, str, str]]],
fields_in_entities: dict[MetricEntity, list[MetricFieldKey]],
intervals: list[datetime],
results,
use_case_id: UseCaseID,
Expand All @@ -1273,15 +1300,15 @@ def __init__(
}

# This is a set of all the `(op, metric_mri, alias)` combinations passed in the metrics_query
self._metrics_query_fields_set = {
self._metrics_query_fields_set: set[MetricFieldKey] = {
(field.op, field.metric_mri, field.alias) for field in metrics_query.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
# the constituents necessary to calculate instances of CompositeEntityDerivedMetric but
# are not necessarily requested in the query definition
self._fields_in_entities_set = {
self._fields_in_entities_set: set[MetricFieldKey] = {
elem for fields_in_entity in fields_in_entities.values() for elem in fields_in_entity
}
self._set_of_constituent_queries = self._fields_in_entities_set.union(
Expand Down Expand Up @@ -1436,7 +1463,7 @@ def resolve_tag_value(value: int | str | None) -> str | None:
except KeyError:
params = None
totals[alias] = metric_obj.run_post_query_function(
totals, params=params, alias=alias
totals, params=_normalize_metric_operation_params(params), alias=alias
)

if series is not None:
Expand All @@ -1451,7 +1478,10 @@ def resolve_tag_value(value: int | str | None) -> str | None:
except KeyError:
params = None
series[alias][idx] = metric_obj.run_post_query_function(
series, params=params, idx=idx, alias=alias
series,
params=_normalize_metric_operation_params(params),
idx=idx,
alias=alias,
)

# Remove the extra fields added due to the constituent metrics that were added
Expand Down
Loading