From 9a816c938f97485e30a56b119cf1753dab866c51 Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Thu, 30 Jan 2025 15:02:29 -0800 Subject: [PATCH 1/9] init wip --- requirements.txt | 2 +- .../R_eap_spans/resolver_time_series.py | 4 + .../R_eap_spans/resolver_trace_item_table.py | 122 +++++++++++------- 3 files changed, 80 insertions(+), 48 deletions(-) diff --git a/requirements.txt b/requirements.txt index daecffff40..94f53a82fb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -29,7 +29,7 @@ python-rapidjson==1.8 redis==4.5.4 sentry-arroyo==2.19.12 sentry-kafka-schemas==0.1.129 -sentry-protos==0.1.57 +sentry-protos==0.1.58 sentry-redis-tools==0.3.0 sentry-relay==0.9.5 sentry-sdk==2.18.0 diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py index 9312a4d7c1..3dcbd11ad7 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py @@ -37,6 +37,7 @@ extract_response_meta, setup_trace_query_settings, ) +from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException from snuba.web.rpc.v1.resolvers import ResolverTimeSeries from snuba.web.rpc.v1.resolvers.R_eap_spans.common.aggregation import ( ExtrapolationContext, @@ -297,6 +298,9 @@ def trace_item_type(cls) -> TraceItemType.ValueType: return TraceItemType.TRACE_ITEM_TYPE_SPAN def resolve(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse: + if in_msg.HasField("expressions"): + raise BadSnubaRPCRequestException("expressions field not yet implemented") + snuba_request = _build_snuba_request(in_msg) res = run_query( dataset=PluggableDataset(name="eap", all_entities=[]), diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py index 5408e56075..91b943c34e 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py @@ -7,6 +7,7 @@ from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import ( AggregationComparisonFilter, AggregationFilter, + Column, TraceItemColumnValues, TraceItemTableRequest, TraceItemTableResponse, @@ -128,6 +129,72 @@ def _convert_order_by( return res +def _get_reliability_context_columns(column: Column) -> list[SelectedExpression]: + """ + extrapolated aggregates need to request extra columns to calculate the reliability of the result. + this function returns the list of columns that need to be requested. + """ + if not column.HasField("aggregation"): + return [] + + if ( + column.aggregation.extrapolation_mode + == ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED + ): + context_columns = [] + confidence_interval_column = get_confidence_interval_column(column.aggregation) + if confidence_interval_column is not None: + context_columns.append( + SelectedExpression( + name=confidence_interval_column.alias, + expression=confidence_interval_column, + ) + ) + + average_sample_rate_column = get_average_sample_rate_column(column.aggregation) + count_column = get_count_column(column.aggregation) + context_columns.append( + SelectedExpression( + name=average_sample_rate_column.alias, + expression=average_sample_rate_column, + ) + ) + context_columns.append( + SelectedExpression(name=count_column.alias, expression=count_column) + ) + return context_columns + return [] + + +def _column_to_expression(column: Column) -> Expression: + """ + Given a column protobuf object, translates it into a Expression object and returns it. + """ + if column.HasField("key"): + return attribute_key_to_expression(column.key) + elif column.HasField("aggregation"): + function_expr = aggregation_to_expression(column.aggregation) + # aggregation label may not be set and the column label takes priority anyways. + function_expr = replace(function_expr, alias=column.label) + return function_expr + elif column.HasField("formula"): + op_to_expr = { + Column.BinaryFormula.OP_ADD: f.add, + Column.BinaryFormula.OP_SUBTRACT: f.subtract, + Column.BinaryFormula.OP_MULTIPLY: f.multiply, + Column.BinaryFormula.OP_DIVIDE: f.divide, + } + return op_to_expr[column.formula.op]( + _column_to_expression(column.formula.left), + _column_to_expression(column.formula.right), + ) + + else: + raise BadSnubaRPCRequestException( + "Column is not one of: aggregate, attribute key, or formula" + ) + + def _build_query(request: TraceItemTableRequest) -> Query: # TODO: This is hardcoded still entity = Entity( @@ -138,54 +205,15 @@ def _build_query(request: TraceItemTableRequest) -> Query: selected_columns = [] for column in request.columns: - if column.HasField("key"): - key_col = attribute_key_to_expression(column.key) - # The key_col expression alias may differ from the column label. That is okay - # the attribute key name is used in the groupby, the column label is just the name of - # the returned attribute value - selected_columns.append( - SelectedExpression(name=column.label, expression=key_col) - ) - elif column.HasField("aggregation"): - function_expr = aggregation_to_expression(column.aggregation) - # aggregation label may not be set and the column label takes priority anyways. - function_expr = replace(function_expr, alias=column.label) - selected_columns.append( - SelectedExpression(name=column.label, expression=function_expr) - ) - - if ( - column.aggregation.extrapolation_mode - == ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED - ): - confidence_interval_column = get_confidence_interval_column( - column.aggregation - ) - if confidence_interval_column is not None: - selected_columns.append( - SelectedExpression( - name=confidence_interval_column.alias, - expression=confidence_interval_column, - ) - ) - - average_sample_rate_column = get_average_sample_rate_column( - column.aggregation - ) - count_column = get_count_column(column.aggregation) - selected_columns.append( - SelectedExpression( - name=average_sample_rate_column.alias, - expression=average_sample_rate_column, - ) - ) - selected_columns.append( - SelectedExpression(name=count_column.alias, expression=count_column) - ) - else: - raise BadSnubaRPCRequestException( - "Column is neither an aggregate or an attribute" + # The key_col expression alias may differ from the column label. That is okay + # the attribute key name is used in the groupby, the column label is just the name of + # the returned attribute value + selected_columns.append( + SelectedExpression( + name=column.label, expression=_column_to_expression(column) ) + ) + selected_columns.extend(_get_reliability_context_columns(column)) res = Query( from_clause=entity, From bb00427e9ddf461e449bddc0a513221cefa25333 Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Thu, 30 Jan 2025 15:15:13 -0800 Subject: [PATCH 2/9] fix mypy --- snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py index 3dcbd11ad7..88a01321d4 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py @@ -298,7 +298,7 @@ def trace_item_type(cls) -> TraceItemType.ValueType: return TraceItemType.TRACE_ITEM_TYPE_SPAN def resolve(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse: - if in_msg.HasField("expressions"): + if len(in_msg.expressions) > 0: raise BadSnubaRPCRequestException("expressions field not yet implemented") snuba_request = _build_snuba_request(in_msg) From b448332aead1913378b6314e25639e8fab57b6fd Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Fri, 31 Jan 2025 13:49:27 -0800 Subject: [PATCH 3/9] initial implementation, might add more test --- .../R_eap_spans/resolver_trace_item_table.py | 5 +- .../test_endpoint_trace_item_table.py | 64 +++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py index 91b943c34e..dadb238638 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py @@ -187,6 +187,7 @@ def _column_to_expression(column: Column) -> Expression: return op_to_expr[column.formula.op]( _column_to_expression(column.formula.left), _column_to_expression(column.formula.right), + alias=column.label, ) else: @@ -283,9 +284,11 @@ def _convert_results( converters[column.label] = lambda x: AttributeValue(val_double=float(x)) elif column.HasField("aggregation"): converters[column.label] = lambda x: AttributeValue(val_double=float(x)) + elif column.HasField("formula"): + converters[column.label] = lambda x: AttributeValue(val_double=float(x)) else: raise BadSnubaRPCRequestException( - "column is neither an attribute or aggregation" + "column is not one of: attribute, aggregation, or formula" ) res: defaultdict[str, TraceItemColumnValues] = defaultdict(TraceItemColumnValues) 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 index b223f6c324..a24d66ac24 100644 --- 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 @@ -2307,6 +2307,70 @@ def test_sparse_aggregate(self, setup_teardown: Any) -> None: ), ] + def test_formula(self, setup_teardown: Any) -> None: + span_ts = BASE_TIME - timedelta(minutes=1) + write_eap_span(span_ts, {"kyles_measurement": 6}, 10) + write_eap_span(span_ts, {"kyles_measurement": 7}, 2) + + ts = Timestamp(seconds=int(BASE_TIME.timestamp())) + hour_ago = int((BASE_TIME - timedelta(hours=1)).timestamp()) + message = TraceItemTableRequest( + meta=RequestMeta( + project_ids=[1, 2, 3], + organization_id=1, + cogs_category="something", + referrer="something", + start_timestamp=Timestamp(seconds=hour_ago), + end_timestamp=ts, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + ), + filter=TraceItemFilter( + exists_filter=ExistsFilter( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="kyles_measurement" + ) + ) + ), + columns=[ + Column( + formula=Column.BinaryFormula( + op=Column.BinaryFormula.OP_DIVIDE, + left=Column( + aggregation=AttributeAggregation( + aggregate=Function.FUNCTION_SUM, + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, + name="kyles_measurement", + ), + ), + label="sum(kyles_measurement)", + ), + right=Column( + aggregation=AttributeAggregation( + aggregate=Function.FUNCTION_COUNT, + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, + name="kyles_measurement", + ), + ), + label="count(kyles_measurement)", + ), + ), + label="sum(kyles_measurement) / count(kyles_measurement)", + ), + ], + limit=1, + ) + response = EndpointTraceItemTable().execute(message) + assert response.column_values == [ + TraceItemColumnValues( + attribute_name="sum(kyles_measurement) / count(kyles_measurement)", + results=[ + AttributeValue(val_double=(74 / 12)), + ], + ), + ] + class TestUtils: def test_apply_labels_to_columns_backward_compat(self) -> None: From da0356c7eb489fef1fb9d51e55b8b1bd0ef14e3b Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Fri, 31 Jan 2025 14:17:28 -0800 Subject: [PATCH 4/9] extrapolation test --- ...endpoint_trace_item_table_extrapolation.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py index 7129f2b781..0b5f16a14c 100644 --- a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py +++ b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py @@ -7,22 +7,28 @@ from google.protobuf.timestamp_pb2 import Timestamp from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import ( Column, + TraceItemColumnValues, TraceItemTableRequest, ) from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta, TraceItemType from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ( AttributeAggregation, AttributeKey, + AttributeValue, ExtrapolationMode, Function, Reliability, ) +from sentry_protos.snuba.v1.trace_item_filter_pb2 import ExistsFilter, TraceItemFilter from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey from snuba.web.rpc.v1.endpoint_trace_item_table import EndpointTraceItemTable from tests.base import BaseApiTest from tests.helpers import write_raw_unprocessed_events +from tests.web.rpc.v1.test_endpoint_trace_item_table.test_endpoint_trace_item_table import ( + write_eap_span, +) _RELEASE_TAG = "backend@24.7.0.dev0+c45b49caed1e5fcbf70097ab3f434b487c359b6b" _SERVER_NAME = "D23CXQ4GK2.local" @@ -708,3 +714,73 @@ def test_count_reliability_with_group_by(self) -> None: assert measurement_reliabilities == [ Reliability.RELIABILITY_LOW, ] # low reliability due to low sample count + + def test_formula(self) -> None: + """ + This test ensures that formulas work with extrapolation. + Reliabilities will not be returned. + """ + span_ts = BASE_TIME - timedelta(minutes=1) + write_eap_span(span_ts, {"kyles_measurement": 6, "server_sample_rate": 0.5}, 10) + write_eap_span(span_ts, {"kyles_measurement": 7, "server_sample_rate": 0.2}, 2) + + ts = Timestamp(seconds=int(BASE_TIME.timestamp())) + hour_ago = int((BASE_TIME - timedelta(hours=1)).timestamp()) + message = TraceItemTableRequest( + meta=RequestMeta( + project_ids=[1, 2, 3], + organization_id=1, + cogs_category="something", + referrer="something", + start_timestamp=Timestamp(seconds=hour_ago), + end_timestamp=ts, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + ), + filter=TraceItemFilter( + exists_filter=ExistsFilter( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="kyles_measurement" + ) + ) + ), + columns=[ + Column( + formula=Column.BinaryFormula( + op=Column.BinaryFormula.OP_DIVIDE, + left=Column( + aggregation=AttributeAggregation( + aggregate=Function.FUNCTION_SUM, + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, + name="kyles_measurement", + ), + extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, + ), + label="sum(kyles_measurement)", + ), + right=Column( + aggregation=AttributeAggregation( + aggregate=Function.FUNCTION_COUNT, + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, + name="kyles_measurement", + ), + extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, + ), + label="count(kyles_measurement)", + ), + ), + label="sum(kyles_measurement) / count(kyles_measurement)", + ), + ], + limit=1, + ) + response = EndpointTraceItemTable().execute(message) + assert response.column_values == [ + TraceItemColumnValues( + attribute_name="sum(kyles_measurement) / count(kyles_measurement)", + results=[ + AttributeValue(val_double=(190 / 30)), + ], + ), + ] From 1f1b38a60294765cad68653745e016c1e3b51689 Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Fri, 31 Jan 2025 14:25:09 -0800 Subject: [PATCH 5/9] update extrapolation test --- .../test_endpoint_trace_item_table_extrapolation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py index 0b5f16a14c..8499e67ded 100644 --- a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py +++ b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py @@ -722,7 +722,7 @@ def test_formula(self) -> None: """ span_ts = BASE_TIME - timedelta(minutes=1) write_eap_span(span_ts, {"kyles_measurement": 6, "server_sample_rate": 0.5}, 10) - write_eap_span(span_ts, {"kyles_measurement": 7, "server_sample_rate": 0.2}, 2) + write_eap_span(span_ts, {"kyles_measurement": 7}, 2) ts = Timestamp(seconds=int(BASE_TIME.timestamp())) hour_ago = int((BASE_TIME - timedelta(hours=1)).timestamp()) @@ -780,7 +780,7 @@ def test_formula(self) -> None: TraceItemColumnValues( attribute_name="sum(kyles_measurement) / count(kyles_measurement)", results=[ - AttributeValue(val_double=(190 / 30)), + AttributeValue(val_double=(134 / 22)), ], ), ] From c3c976206d5f8caae724b48893abe98d48ba4817 Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Fri, 31 Jan 2025 15:01:54 -0800 Subject: [PATCH 6/9] test non-agg-formula --- .../R_eap_spans/resolver_trace_item_table.py | 35 ++++--- .../test_endpoint_trace_item_table.py | 93 ++++++++++++++++++- 2 files changed, 115 insertions(+), 13 deletions(-) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py index dadb238638..912ec45e61 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py @@ -126,6 +126,13 @@ def _convert_order_by( expression=aggregation_to_expression(x.column.aggregation), ) ) + elif x.column.HasField("formula"): + res.append( + OrderBy( + direction=direction, + expression=_formula_to_expression(x.column.formula), + ) + ) return res @@ -166,6 +173,19 @@ def _get_reliability_context_columns(column: Column) -> list[SelectedExpression] return [] +def _formula_to_expression(formula: Column.BinaryFormula) -> Expression: + op_to_expr = { + Column.BinaryFormula.OP_ADD: f.plus, + Column.BinaryFormula.OP_SUBTRACT: f.minus, + Column.BinaryFormula.OP_MULTIPLY: f.multiply, + Column.BinaryFormula.OP_DIVIDE: f.divide, + } + return op_to_expr[formula.op]( + _column_to_expression(formula.left), + _column_to_expression(formula.right), + ) + + def _column_to_expression(column: Column) -> Expression: """ Given a column protobuf object, translates it into a Expression object and returns it. @@ -178,18 +198,9 @@ def _column_to_expression(column: Column) -> Expression: function_expr = replace(function_expr, alias=column.label) return function_expr elif column.HasField("formula"): - op_to_expr = { - Column.BinaryFormula.OP_ADD: f.add, - Column.BinaryFormula.OP_SUBTRACT: f.subtract, - Column.BinaryFormula.OP_MULTIPLY: f.multiply, - Column.BinaryFormula.OP_DIVIDE: f.divide, - } - return op_to_expr[column.formula.op]( - _column_to_expression(column.formula.left), - _column_to_expression(column.formula.right), - alias=column.label, - ) - + formula_expr = _formula_to_expression(column.formula) + formula_expr = replace(formula_expr, alias=column.label) + return formula_expr else: raise BadSnubaRPCRequestException( "Column is not one of: aggregate, attribute key, or formula" 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 index a24d66ac24..6e3a83e546 100644 --- 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 @@ -2307,7 +2307,11 @@ def test_sparse_aggregate(self, setup_teardown: Any) -> None: ), ] - def test_formula(self, setup_teardown: Any) -> None: + def test_agg_formula(self, setup_teardown: Any) -> None: + """ + ensures formulas of aggregates work + ex sum(my_attribute) / count(my_attribute) + """ span_ts = BASE_TIME - timedelta(minutes=1) write_eap_span(span_ts, {"kyles_measurement": 6}, 10) write_eap_span(span_ts, {"kyles_measurement": 7}, 2) @@ -2371,6 +2375,93 @@ def test_formula(self, setup_teardown: Any) -> None: ), ] + def test_non_agg_formula(self, setup_teardown: Any) -> None: + """ + ensures formulas of non-aggregates work + ex: my_attribute + my_other_attribute + """ + span_ts = BASE_TIME - timedelta(minutes=1) + write_eap_span(span_ts, {"kyles_measurement": -1, "my_other_attribute": 1}, 4) + write_eap_span(span_ts, {"kyles_measurement": 3, "my_other_attribute": 2}, 2) + write_eap_span(span_ts, {"kyles_measurement": 10, "my_other_attribute": 3}, 1) + + ts = Timestamp(seconds=int(BASE_TIME.timestamp())) + hour_ago = int((BASE_TIME - timedelta(hours=1)).timestamp()) + message = TraceItemTableRequest( + meta=RequestMeta( + project_ids=[1, 2, 3], + organization_id=1, + cogs_category="something", + referrer="something", + start_timestamp=Timestamp(seconds=hour_ago), + end_timestamp=ts, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + ), + filter=TraceItemFilter( + exists_filter=ExistsFilter( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="kyles_measurement" + ) + ) + ), + columns=[ + Column( + formula=Column.BinaryFormula( + op=Column.BinaryFormula.OP_ADD, + left=Column( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="kyles_measurement" + ) + ), + right=Column( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="my_other_attribute" + ) + ), + ), + label="kyles_measurement + my_other_attribute", + ), + ], + order_by=[ + TraceItemTableRequest.OrderBy( + column=Column( + formula=Column.BinaryFormula( + op=Column.BinaryFormula.OP_ADD, + left=Column( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, + name="kyles_measurement", + ) + ), + right=Column( + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, + name="my_other_attribute", + ) + ), + ), + label="kyles_measurement + my_other_attribute", + ) + ), + ], + limit=50, + ) + response = EndpointTraceItemTable().execute(message) + assert response.column_values == [ + TraceItemColumnValues( + attribute_name="kyles_measurement + my_other_attribute", + results=[ + AttributeValue(val_double=0), + AttributeValue(val_double=0), + AttributeValue(val_double=0), + AttributeValue(val_double=0), + AttributeValue(val_double=5), + AttributeValue(val_double=5), + AttributeValue(val_double=13), + ], + ), + ] + class TestUtils: def test_apply_labels_to_columns_backward_compat(self) -> None: From 42925021e89e1850a468dcb1b6c52ebc8d3acd33 Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Mon, 3 Feb 2025 11:38:24 -0800 Subject: [PATCH 7/9] error msg for logs and uptime --- .../rpc/v1/resolvers/R_ourlogs/resolver_trace_item_table.py | 2 +- .../v1/resolvers/R_uptime_checks/resolver_trace_item_table.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/snuba/web/rpc/v1/resolvers/R_ourlogs/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_ourlogs/resolver_trace_item_table.py index c44ae6aaf4..1f8d32049d 100644 --- a/snuba/web/rpc/v1/resolvers/R_ourlogs/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_ourlogs/resolver_trace_item_table.py @@ -78,7 +78,7 @@ def _build_query(request: TraceItemTableRequest) -> Query: ) else: raise BadSnubaRPCRequestException( - "requested attribute is not a column (aggregation not supported for logs)" + "requested attribute is not a column (aggregation and formulanot supported for logs)" ) res = Query( diff --git a/snuba/web/rpc/v1/resolvers/R_uptime_checks/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_uptime_checks/resolver_trace_item_table.py index 1e541b0ea0..d8aa2f42d9 100644 --- a/snuba/web/rpc/v1/resolvers/R_uptime_checks/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_uptime_checks/resolver_trace_item_table.py @@ -144,6 +144,10 @@ def _build_query(request: TraceItemTableRequest) -> Query: selected_columns.append( SelectedExpression(name=column.label, expression=function_expr) ) + elif column.HasField("formula"): + raise BadSnubaRPCRequestException( + "formulas are not supported for uptime checks" + ) else: raise BadSnubaRPCRequestException( "Column is neither an aggregate or an attribute" From 6d91e4ed7e360e2ee8eccbf350e9fe24381efa2b Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Mon, 3 Feb 2025 17:05:12 -0800 Subject: [PATCH 8/9] pr feedback --- .../R_eap_spans/resolver_trace_item_table.py | 29 +++++++++---------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py index 912ec45e61..01fc41398e 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py @@ -27,7 +27,7 @@ from snuba.query import OrderBy, OrderByDirection, SelectedExpression from snuba.query.data_source.simple import Entity from snuba.query.dsl import Functions as f -from snuba.query.dsl import and_cond, or_cond +from snuba.query.dsl import and_cond, literal, or_cond from snuba.query.expressions import Expression from snuba.query.logical import Query from snuba.query.query_settings import HTTPQuerySettings @@ -56,6 +56,13 @@ _DEFAULT_ROW_LIMIT = 10_000 +OP_TO_EXPR = { + Column.BinaryFormula.OP_ADD: f.plus, + Column.BinaryFormula.OP_SUBTRACT: f.minus, + Column.BinaryFormula.OP_MULTIPLY: f.multiply, + Column.BinaryFormula.OP_DIVIDE: f.divide, +} + def aggregation_filter_to_expression(agg_filter: AggregationFilter) -> Expression: op_to_expr = { @@ -130,7 +137,7 @@ def _convert_order_by( res.append( OrderBy( direction=direction, - expression=_formula_to_expression(x.column.formula), + expression=literal(x.column.label), ) ) return res @@ -173,19 +180,6 @@ def _get_reliability_context_columns(column: Column) -> list[SelectedExpression] return [] -def _formula_to_expression(formula: Column.BinaryFormula) -> Expression: - op_to_expr = { - Column.BinaryFormula.OP_ADD: f.plus, - Column.BinaryFormula.OP_SUBTRACT: f.minus, - Column.BinaryFormula.OP_MULTIPLY: f.multiply, - Column.BinaryFormula.OP_DIVIDE: f.divide, - } - return op_to_expr[formula.op]( - _column_to_expression(formula.left), - _column_to_expression(formula.right), - ) - - def _column_to_expression(column: Column) -> Expression: """ Given a column protobuf object, translates it into a Expression object and returns it. @@ -198,7 +192,10 @@ def _column_to_expression(column: Column) -> Expression: function_expr = replace(function_expr, alias=column.label) return function_expr elif column.HasField("formula"): - formula_expr = _formula_to_expression(column.formula) + formula_expr = OP_TO_EXPR[column.formula.op]( + _column_to_expression(column.formula.left), + _column_to_expression(column.formula.right), + ) formula_expr = replace(formula_expr, alias=column.label) return formula_expr else: From a53e0db769deaac50bb34b47dd52feb31f39ca44 Mon Sep 17 00:00:00 2001 From: Kyle Mumma Date: Tue, 4 Feb 2025 12:06:03 -0800 Subject: [PATCH 9/9] fix pr feedback --- .../R_eap_spans/resolver_trace_item_table.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py index 01fc41398e..87635b7c24 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py @@ -27,7 +27,7 @@ from snuba.query import OrderBy, OrderByDirection, SelectedExpression from snuba.query.data_source.simple import Entity from snuba.query.dsl import Functions as f -from snuba.query.dsl import and_cond, literal, or_cond +from snuba.query.dsl import and_cond, or_cond from snuba.query.expressions import Expression from snuba.query.logical import Query from snuba.query.query_settings import HTTPQuerySettings @@ -137,7 +137,7 @@ def _convert_order_by( res.append( OrderBy( direction=direction, - expression=literal(x.column.label), + expression=_formula_to_expression(x.column.formula), ) ) return res @@ -180,6 +180,13 @@ def _get_reliability_context_columns(column: Column) -> list[SelectedExpression] return [] +def _formula_to_expression(formula: Column.BinaryFormula) -> Expression: + return OP_TO_EXPR[formula.op]( + _column_to_expression(formula.left), + _column_to_expression(formula.right), + ) + + def _column_to_expression(column: Column) -> Expression: """ Given a column protobuf object, translates it into a Expression object and returns it. @@ -192,10 +199,7 @@ def _column_to_expression(column: Column) -> Expression: function_expr = replace(function_expr, alias=column.label) return function_expr elif column.HasField("formula"): - formula_expr = OP_TO_EXPR[column.formula.op]( - _column_to_expression(column.formula.left), - _column_to_expression(column.formula.right), - ) + formula_expr = _formula_to_expression(column.formula) formula_expr = replace(formula_expr, alias=column.label) return formula_expr else: