Browse Source

feat(metrics): extend counter condtion

Ogi 3 days ago
parent
commit
bafc23d7e0

+ 40 - 8
src/sentry/snuba/metrics/span_attribute_extraction.py

@@ -1,8 +1,14 @@
 from collections.abc import Sequence
-from typing import Any, Literal, NotRequired, TypedDict
+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,
+    SearchKey,
+    SearchValue,
+)
 from sentry.sentry_metrics.extraction_rules import MetricsExtractionRule
 from sentry.snuba.metrics.extraction import RuleCondition, SearchQueryConverter, TagSpec
 
@@ -85,20 +91,20 @@ def convert_to_metric_spec(extraction_rule: MetricsExtractionRule) -> SpanAttrib
         "category": "span",
         "mri": extraction_rule.generate_mri(),
         "field": field,
-        "tags": _get_tags(extraction_rule.tags, parsed_conditions),
-        "condition": _get_rule_condition(parsed_conditions),
+        "tags": _get_tags(extraction_rule, parsed_conditions),
+        "condition": _get_rule_condition(extraction_rule, parsed_conditions),
     }
 
 
 def _get_field(extraction_rule: MetricsExtractionRule) -> str | None:
-    if extraction_rule.type == "c":
+    if _is_counter(extraction_rule):
         return None
 
     return _map_span_attribute_name(extraction_rule.span_attribute)
 
 
 def _get_tags(
-    explicitly_defined_tags: set[str], conditions: Sequence[QueryToken] | None
+    extraction_rule: MetricsExtractionRule, conditions: Sequence[QueryToken] | None
 ) -> list[TagSpec]:
     """
     Merges the explicitly defined tags with the tags extracted from the search conditions.
@@ -106,7 +112,7 @@ def _get_tags(
     token_list = _flatten_query_tokens(conditions) if conditions else []
     search_token_keys = {token.key.name for token in token_list}
 
-    tag_keys = explicitly_defined_tags.union(search_token_keys)
+    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)]
 
@@ -133,7 +139,16 @@ def _parse_conditions(conditions: Sequence[str] | None) -> Sequence[QueryToken]:
     return event_search.parse_search_query(search_query)
 
 
-def _get_rule_condition(parsed_search_query: Sequence[Any] | None) -> RuleCondition | None:
+def _get_rule_condition(
+    extraction_rule: MetricsExtractionRule, parsed_conditions: Sequence[QueryToken]
+) -> RuleCondition | None:
+    if _is_counter(extraction_rule):
+        parsed_search_query = _append_exists_condition(
+            parsed_conditions, extraction_rule.span_attribute
+        )
+    else:
+        parsed_search_query = parsed_conditions
+
     if not parsed_search_query:
         return None
 
@@ -142,6 +157,19 @@ def _get_rule_condition(parsed_search_query: Sequence[Any] | None) -> RuleCondit
     ).convert()
 
 
+def _append_exists_condition(
+    parsed_conditions: Sequence[QueryToken], span_attribute: str
+) -> Sequence[QueryToken]:
+    exists_search_filter = SearchFilter(
+        key=SearchKey("has"), operator="!=", value=SearchValue(span_attribute)
+    )
+
+    if not parsed_conditions:
+        return [exists_search_filter]
+
+    return [ParenExpression(children=[*parsed_conditions, exists_search_filter])]
+
+
 def _map_span_attribute_name(span_attribute: str) -> str:
     if span_attribute in _TOP_LEVEL_SPAN_ATTRIBUTES:
         return span_attribute
@@ -154,3 +182,7 @@ def _map_span_attribute_name(span_attribute: str) -> str:
     sanitized_span_attr = span_attribute.replace(".", "\\.")
 
     return f"{prefix}.{sanitized_span_attr}"
+
+
+def _is_counter(extraction_rule: MetricsExtractionRule) -> bool:
+    return extraction_rule.type == "c"

+ 23 - 2
tests/sentry/snuba/metrics/test_span_attribute_extraction.py

@@ -154,12 +154,33 @@ def test_counter():
 
     assert not metric_spec["field"]
     assert metric_spec["mri"] == "c:custom/foobar@none"
-    assert metric_spec["tags"] == []
+    assert metric_spec["condition"] == {
+        "inner": {"name": "span.data.has", "op": "eq", "value": "foobar"},
+        "op": "not",
+    }
+
+
+def test_counter_extends_conditions():
+    rule = MetricsExtractionRule(
+        span_attribute="foobar", type="c", unit="none", tags=set(), conditions=["abc:xyz"]
+    )
+
+    metric_spec = convert_to_metric_spec(rule)
+
+    assert not metric_spec["field"]
+    assert metric_spec["mri"] == "c:custom/foobar@none"
+    assert metric_spec["condition"] == {
+        "op": "and",
+        "inner": [
+            {"op": "eq", "name": "span.data.abc", "value": "xyz"},
+            {"inner": {"name": "span.data.has", "op": "eq", "value": "foobar"}, "op": "not"},
+        ],
+    }
 
 
 def test_empty_conditions():
     rule = MetricsExtractionRule(
-        span_attribute="foobar", type="c", unit="none", tags=set(), conditions=[""]
+        span_attribute="foobar", type="d", unit="none", tags=set(), conditions=[""]
     )
 
     metric_spec = convert_to_metric_spec(rule)