clean up

getsentry · Feb 14, 2025 · 462a2a7 · 462a2a7
1 parent 24b3a30
commit 462a2a7
Showing 1 changed file with 16 additions and 18 deletions.
diff --git a/snuba/web/rpc/v1/endpoint_trace_item_table.py b/snuba/web/rpc/v1/endpoint_trace_item_table.py
@@ -5,6 +5,7 @@
     AttributeConditionalAggregation,
 )
 from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import (
+    AggregationComparisonFilter,
     AggregationFilter,
     Column,
     TraceItemTableRequest,
@@ -87,18 +88,23 @@ def _transform_request(request: TraceItemTableRequest) -> TraceItemTableRequest:
 
 
 def convert_to_conditional_aggregation(in_msg: TraceItemTableRequest) -> None:
+    def _add_conditional_aggregation(
+        input: Column | AggregationComparisonFilter,
+    ) -> None:
+        aggregation = input.aggregation
+        input.conditional_aggregation.CopyFrom(
+            AttributeConditionalAggregation(
+                aggregate=aggregation.aggregate,
+                key=aggregation.key,
+                label=aggregation.label,
+                extrapolation_mode=aggregation.extrapolation_mode,
+            )
+        )
+
     def _convert(input: Column | AggregationFilter) -> None:
         if isinstance(input, Column):
             if input.HasField("aggregation"):
-                aggregation = input.aggregation
-                input.conditional_aggregation.CopyFrom(
-                    AttributeConditionalAggregation(
-                        aggregate=aggregation.aggregate,
-                        key=aggregation.key,
-                        label=aggregation.label,
-                        extrapolation_mode=aggregation.extrapolation_mode,
-                    )
-                )
+                _add_conditional_aggregation(input)
 
             if input.HasField("formula"):
                 _convert(input.formula.left)
@@ -113,15 +119,7 @@ def _convert(input: Column | AggregationFilter) -> None:
                     _convert(aggregation_filter)
             if input.HasField("comparison_filter"):
                 if input.comparison_filter.HasField("aggregation"):
-                    aggregation = input.comparison_filter.aggregation
-                    input.comparison_filter.conditional_aggregation.CopyFrom(
-                        AttributeConditionalAggregation(
-                            aggregate=aggregation.aggregate,
-                            key=aggregation.key,
-                            label=aggregation.label,
-                            extrapolation_mode=aggregation.extrapolation_mode,
-                        )
-                    )
+                    _add_conditional_aggregation(input.comparison_filter)
 
     for column in in_msg.columns:
         _convert(column)