From c86f98c44ffd5ee5b301b87d961f1d7405776826 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 2 Aug 2026 12:07:28 +0000 Subject: [PATCH 1/2] typing: strongly type metrics query_builder module Co-authored-by: Armen Zambrano G. --- pyproject.toml | 1 - src/sentry/snuba/metrics/query_builder.py | 157 +++++++++++++--------- 2 files changed, 96 insertions(+), 62 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e499228a2f0d..3b790a8ca133 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = [ diff --git a/src/sentry/snuba/metrics/query_builder.py b/src/sentry/snuba/metrics/query_builder.py index 066d553cc6f3..98379b729671 100644 --- a/src/sentry/snuba/metrics/query_builder.py +++ b/src/sentry/snuba/metrics/query_builder.py @@ -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 ( @@ -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, @@ -71,6 +72,8 @@ DerivedMetricParseException, MetricDoesNotExistException, MetricEntity, + MetricOperationType, + OPERATIONS, get_num_intervals, get_timestamp_column_name, require_rhs_condition_resolution, @@ -93,6 +96,20 @@ 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: + if params is None: + return None + + return { + key: value + for key, value in params.items() + if value is not None and isinstance(value, (str, int, float)) + } or None def _strip_project_id(condition: Condition) -> Condition | None: @@ -134,7 +151,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): @@ -747,11 +767,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: @@ -823,82 +846,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]: @@ -920,21 +938,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: @@ -1069,6 +1095,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 ) @@ -1099,11 +1127,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 @@ -1117,19 +1145,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) @@ -1201,7 +1233,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 @@ -1255,7 +1287,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, @@ -1273,7 +1305,7 @@ 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 @@ -1281,7 +1313,7 @@ def __init__( # 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( @@ -1436,7 +1468,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: @@ -1451,7 +1483,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 From 18cf7cb966b4917149bf468199d64e5b87f0f3ca Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 2 Aug 2026 12:11:03 +0000 Subject: [PATCH 2/2] typing: preserve metric operation params behavior Co-authored-by: Armen Zambrano G. --- src/sentry/snuba/metrics/query_builder.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/sentry/snuba/metrics/query_builder.py b/src/sentry/snuba/metrics/query_builder.py index 98379b729671..b3d019f1a145 100644 --- a/src/sentry/snuba/metrics/query_builder.py +++ b/src/sentry/snuba/metrics/query_builder.py @@ -102,14 +102,9 @@ def _normalize_metric_operation_params( params: dict[str, None | str | int | float | Sequence[tuple[str | int, ...]]] | None, ) -> MetricOperationParams | None: - if params is None: - return None - - return { - key: value - for key, value in params.items() - if value is not None and isinstance(value, (str, int, float)) - } or 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: