feat(metrics): numeric span attribute extraction (#74190)

Author: obostjancic

src/sentry/snuba/metrics/span_attribute_extraction.py

@@ -2,7 +2,7 @@
 from typing import Literal, NotRequired, TypedDict
 
 from sentry.api import event_search
-from sentry.api.event_search import ParenExpression, QueryToken, SearchFilter
+from sentry.api.event_search import ParenExpression, QueryToken, SearchFilter, SearchValue
 from sentry.relay.types import RuleCondition
 from sentry.sentry_metrics.extraction_rules import MetricsExtractionRule
 from sentry.snuba.metrics.extraction import SearchQueryConverter, TagSpec
@@ -70,6 +70,10 @@
 
 
 class SpanAttributeMetricSpec(TypedDict):
+    """
+    Represents a metric extraction rule to that extracts metrics from span attributes.
+    """
+
     category: Literal["span"]
     mri: str
     field: NotRequired[str | None]
@@ -78,46 +82,50 @@ class SpanAttributeMetricSpec(TypedDict):
 
 
 def convert_to_metric_spec(extraction_rule: MetricsExtractionRule) -> SpanAttributeMetricSpec:
+    """
+    Converts a persisted MetricsExtractionRule into a SpanAttributeMetricSpec that satisfies
+    MetricSpec of relay metric extraction config.
+    """
 
-    field = _get_field(extraction_rule)
-
-    # TODO(metrics): simplify MetricsExtractionRule in a follwup PR
-    parsed_conditions = _parse_conditions([extraction_rule.condition])
+    parsed_search_query = event_search.parse_search_query(extraction_rule.condition)
+    extended_search_query = _extend_search_query(parsed_search_query)
 
     return {
         "category": "span",
         "mri": extraction_rule.generate_mri(),
-        "field": field,
-        "tags": _get_tags(extraction_rule, parsed_conditions),
-        "condition": _get_rule_condition(extraction_rule, parsed_conditions),
+        "field": _get_field(extraction_rule),
+        "tags": _get_tags(extraction_rule, parsed_search_query),
+        "condition": _get_rule_condition(extraction_rule, extended_search_query),
     }
 
 
-def _get_field(extraction_rule: MetricsExtractionRule) -> str | None:
-    if _is_counter(extraction_rule):
-        return None
-
-    return _map_span_attribute_name(extraction_rule.span_attribute)
+# Tag extraction functions
 
 
 def _get_tags(
-    extraction_rule: MetricsExtractionRule, conditions: Sequence[QueryToken] | None
+    extraction_rule: MetricsExtractionRule, parsed_search_query: Sequence[QueryToken] | None
 ) -> list[TagSpec]:
     """
-    Merges the explicitly defined tags with the tags extracted from the search conditions.
+    Merges the explicitly defined tags with the tags extracted from the search query.
     """
-    token_list = _flatten_query_tokens(conditions) if conditions else []
+
+    token_list = _flatten_query_tokens(parsed_search_query) if parsed_search_query else []
     search_token_keys = {token.key.name for token in token_list}
 
     tag_keys = extraction_rule.tags.union(search_token_keys)
 
     return [TagSpec(key=key, field=_map_span_attribute_name(key)) for key in sorted(tag_keys)]
 
 
-def _flatten_query_tokens(conditions: Sequence[QueryToken]) -> list[SearchFilter]:
+def _flatten_query_tokens(parsed_search_query: Sequence[QueryToken]) -> list[SearchFilter]:
+    """
+    Takes a parsed search query and flattens it into a list of SearchFilter tokens.
+    Removes any parenthesis and boolean operators.
+    """
+
     query_tokens: list[SearchFilter] = []
 
-    for token in conditions:
+    for token in parsed_search_query:
         if isinstance(token, SearchFilter):
             query_tokens.append(token)
         elif isinstance(token, ParenExpression):
@@ -126,42 +134,100 @@ def _flatten_query_tokens(conditions: Sequence[QueryToken]) -> list[SearchFilter
     return query_tokens
 
 
-def _parse_conditions(conditions: Sequence[str] | None) -> Sequence[QueryToken]:
-    if not conditions:
-        return []
+# Condition string parsing and transformation functions
 
-    non_empty_conditions = [condition for condition in conditions if condition]
 
-    search_query = " or ".join([f"({condition})" for condition in non_empty_conditions])
-    return event_search.parse_search_query(search_query)
+def _extend_search_query(parsed_search_query: Sequence[QueryToken]) -> Sequence[QueryToken]:
+    return _visit_numeric_tokens(parsed_search_query)
+
+
+def _visit_numeric_tokens(parsed_search_query: Sequence[QueryToken]) -> list[QueryToken]:
+    """
+    Visits each token in the parsed search query and converts numeric tokens into paren expressions.
+    """
+
+    query_tokens: list[QueryToken] = []
+
+    for token in parsed_search_query:
+        if isinstance(token, SearchFilter):
+            query_tokens.append(_extend_numeric_token(token))
+        elif isinstance(token, ParenExpression):
+            query_tokens = query_tokens + _visit_numeric_tokens(token.children)
+        else:
+            query_tokens.append(token)
+
+    return query_tokens
+
+
+def _extend_numeric_token(token: SearchFilter) -> ParenExpression | SearchFilter:
+    """
+    Since all search filter values are parsed as strings by default, we need to make sure that
+    numeric values are treated as such when constructing the rule condition. This function
+    expands the original token into a paren expression if the value is a numeric string.
+
+    Example:
+    `key:'123'` -> `key:'123' OR key:123`
+    `key:['123', '456']` -> `key:['123', '456'] OR key:[123, 456]`
+    """
+
+    if token.operator == "=" or token.operator == "!=":
+        if not str(token.value.value).isdigit():
+            return token
+
+        numeric_value_token = SearchFilter(
+            key=token.key, operator=token.operator, value=SearchValue(int(token.value.value))
+        )
+
+    elif token.is_in_filter:
+        str_values = [str(value) for value in token.value.value]
+        if not all(value.isdigit() for value in str_values):
+            return token
+
+        numeric_values = [int(value) for value in str_values]
+        numeric_value_token = SearchFilter(
+            key=token.key, operator=token.operator, value=SearchValue(numeric_values)
+        )
+
+    return ParenExpression(
+        children=[
+            token,
+            "OR",
+            numeric_value_token,
+        ]
+    )
+
+
+# Conversion to RuleCondition functions
 
 
 def _get_rule_condition(
-    extraction_rule: MetricsExtractionRule, parsed_conditions: Sequence[QueryToken]
+    extraction_rule: MetricsExtractionRule, parsed_search_query: Sequence[QueryToken]
 ) -> RuleCondition | None:
     if _is_counter(extraction_rule):
-        return _get_counter_rule_condition(extraction_rule, parsed_conditions)
+        return _get_counter_rule_condition(extraction_rule, parsed_search_query)
 
-    if not parsed_conditions:
+    if not parsed_search_query:
         return None
 
-    return SearchQueryConverter(parsed_conditions, field_mapper=_map_span_attribute_name).convert()
+    return SearchQueryConverter(
+        parsed_search_query, field_mapper=_map_span_attribute_name
+    ).convert()
 
 
 def _get_counter_rule_condition(
-    extraction_rule: MetricsExtractionRule, parsed_conditions: Sequence[QueryToken]
+    extraction_rule: MetricsExtractionRule, parsed_search_query: Sequence[QueryToken]
 ) -> RuleCondition | None:
     is_top_level = extraction_rule.span_attribute in _TOP_LEVEL_SPAN_ATTRIBUTES
 
-    if not parsed_conditions:
-        # temporary workaround for span.duration counter metric
+    if not parsed_search_query:
+        # workaround for span.duration and other top level attributes that are always present
         if is_top_level:
             return None
 
         return _get_exists_condition(extraction_rule.span_attribute)
 
     condition_dict = SearchQueryConverter(
-        parsed_conditions, field_mapper=_map_span_attribute_name
+        parsed_search_query, field_mapper=_map_span_attribute_name
     ).convert()
 
     if is_top_level:
@@ -194,6 +260,9 @@ def _get_exists_condition(span_attribute: str) -> RuleCondition:
     }
 
 
+# General helpers
+
+
 def _map_span_attribute_name(span_attribute: str) -> str:
     if span_attribute in _TOP_LEVEL_SPAN_ATTRIBUTES:
         return span_attribute
@@ -208,3 +277,10 @@ def _map_span_attribute_name(span_attribute: str) -> str:
 
 def _is_counter(extraction_rule: MetricsExtractionRule) -> bool:
     return extraction_rule.type == "c"
+
+
+def _get_field(extraction_rule: MetricsExtractionRule) -> str | None:
+    if _is_counter(extraction_rule):
+        return None
+
+    return _map_span_attribute_name(extraction_rule.span_attribute)

tests/sentry/snuba/metrics/test_span_attribute_extraction.py

@@ -194,3 +194,64 @@ def test_empty_conditions():
     metric_spec = convert_to_metric_spec(rule)
 
     assert not metric_spec["condition"]
+
+
+def test_numeric_conditions():
+    rule = MetricsExtractionRule(
+        span_attribute="foobar",
+        type="d",
+        unit="none",
+        tags=set(),
+        condition="foo:123 and bar:[123,456]",
+        id=1,
+    )
+
+    metric_spec = convert_to_metric_spec(rule)
+
+    assert metric_spec["condition"] == {
+        "op": "and",
+        "inner": [
+            {
+                "op": "or",
+                "inner": [
+                    {"op": "eq", "name": "span.data.foo", "value": "123"},
+                    {"op": "eq", "name": "span.data.foo", "value": 123},
+                ],
+            },
+            {
+                "op": "or",
+                "inner": [
+                    {"op": "eq", "name": "span.data.bar", "value": ["123", "456"]},
+                    {"op": "eq", "name": "span.data.bar", "value": [123, 456]},
+                ],
+            },
+        ],
+    }
+
+
+def test_mixed_conditions():
+    rule = MetricsExtractionRule(
+        span_attribute="foobar",
+        type="d",
+        unit="none",
+        tags=set(),
+        condition="foo:123 and bar:True and baz:qux",
+        id=1,
+    )
+
+    metric_spec = convert_to_metric_spec(rule)
+
+    assert metric_spec["condition"] == {
+        "op": "and",
+        "inner": [
+            {
+                "op": "or",
+                "inner": [
+                    {"op": "eq", "name": "span.data.foo", "value": "123"},
+                    {"op": "eq", "name": "span.data.foo", "value": 123},
+                ],
+            },
+            {"op": "eq", "name": "span.data.bar", "value": "True"},
+            {"op": "eq", "name": "span.data.baz", "value": "qux"},
+        ],
+    }