more fixes....

getsentry · Feb 14, 2025 · ac24c55 · ac24c55
1 parent 4d57b83
commit ac24c55
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 19 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 (
+    AggregationFilter,
     Column,
     TraceItemTableRequest,
     TraceItemTableResponse,
@@ -86,26 +87,48 @@ def _transform_request(request: TraceItemTableRequest) -> TraceItemTableRequest:
 
 
 def convert_to_conditional_aggregation(in_msg: TraceItemTableRequest) -> None:
-    def _convert(column: Column) -> None:
-        if column.HasField("aggregation"):
-            aggregation = column.aggregation
-            column.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,
+                    )
                 )
-            )
 
-        if column.HasField("formula"):
-            _convert(column.formula.left)
-            _convert(column.formula.right)
+            if input.HasField("formula"):
+                _convert(input.formula.left)
+                _convert(input.formula.right)
+
+        if isinstance(input, AggregationFilter):
+            if input.HasField("and_filter"):
+                for aggregation_filter in input.and_filter.filters:
+                    _convert(aggregation_filter)
+            if input.HasField("or_filter"):
+                for aggregation_filter in input.or_filter.filters:
+                    _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,
+                        )
+                    )
 
     for column in in_msg.columns:
         _convert(column)
     for ob in in_msg.order_by:
         _convert(ob.column)
+    if in_msg.HasField("aggregation_filter"):
+        _convert(in_msg.aggregation_filter)
 
 
 class EndpointTraceItemTable(

diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py
@@ -1916,14 +1916,12 @@ def test_rachel(self) -> None:
             ],
         )
         response = EndpointTraceItemTable().execute(message)
-        print(response.column_values)
         assert response.column_values == [
             TraceItemColumnValues(
                 attribute_name="avg_sample(sampling_rate)",
                 results=[
                     AttributeValue(val_double=0.475),
                 ],
-                reliabilities=[Reliability.RELIABILITY_LOW],
             ),
             TraceItemColumnValues(
                 attribute_name="count()",
@@ -2791,7 +2789,7 @@ def test_apply_labels_to_columns_backward_compat(self) -> None:
         message = TraceItemTableRequest(
             columns=[
                 Column(
-                    aggregation=AttributeAggregation(
+                    conditional_aggregation=AttributeConditionalAggregation(
                         aggregate=Function.FUNCTION_AVG,
                         key=AttributeKey(
                             type=AttributeKey.TYPE_FLOAT, name="custom_measurement"
@@ -2801,7 +2799,7 @@ def test_apply_labels_to_columns_backward_compat(self) -> None:
                     )
                 ),
                 Column(
-                    aggregation=AttributeAggregation(
+                    conditional_aggregation=AttributeConditionalAggregation(
                         aggregate=Function.FUNCTION_AVG,
                         key=AttributeKey(
                             type=AttributeKey.TYPE_FLOAT, name="custom_measurement"
@@ -2823,7 +2821,7 @@ def test_apply_labels_to_columns(self) -> None:
         message = TraceItemTableRequest(
             columns=[
                 Column(
-                    aggregation=AttributeAggregation(
+                    conditional_aggregation=AttributeConditionalAggregation(
                         aggregate=Function.FUNCTION_AVG,
                         key=AttributeKey(
                             type=AttributeKey.TYPE_DOUBLE, name="custom_measurement"
@@ -2833,7 +2831,7 @@ def test_apply_labels_to_columns(self) -> None:
                     )
                 ),
                 Column(
-                    aggregation=AttributeAggregation(
+                    conditional_aggregation=AttributeConditionalAggregation(
                         aggregate=Function.FUNCTION_AVG,
                         key=AttributeKey(
                             type=AttributeKey.TYPE_DOUBLE, name="custom_measurement"