From 4026433627387f9c46d8e39ca5357d07a080bc47 Mon Sep 17 00:00:00 2001 From: Rachel Chen Date: Thu, 12 Dec 2024 13:56:38 -0800 Subject: [PATCH 1/2] c --- snuba/web/rpc/common/common.py | 15 ++++ .../v1/endpoint_trace_item_attribute_names.py | 28 ++---- .../web/rpc/v1/trace_item_attribute_values.py | 88 +++++++++++++------ .../v1/test_trace_item_attribute_values_v1.py | 40 ++++++--- 4 files changed, 113 insertions(+), 58 deletions(-) diff --git a/snuba/web/rpc/common/common.py b/snuba/web/rpc/common/common.py index 923f0c9aa5..ad56136d42 100644 --- a/snuba/web/rpc/common/common.py +++ b/snuba/web/rpc/common/common.py @@ -390,3 +390,18 @@ def base_conditions_and(meta: RequestMeta, *other_exprs: Expression) -> Expressi ), *other_exprs, ) + + +def convert_filter_offset(filter_offset: TraceItemFilter) -> Expression: + if not filter_offset.HasField("comparison_filter"): + raise TypeError("filter_offset needs to be a comparison filter") + if filter_offset.comparison_filter.op != ComparisonFilter.OP_GREATER_THAN: + raise TypeError("filter_offset must use the greater than comparison") + + k_expression = column(filter_offset.comparison_filter.key.name) + v = filter_offset.comparison_filter.value + value_type = v.WhichOneof("value") + if value_type != "val_str": + raise BadSnubaRPCRequestException("please provide a string for filter offset") + + return f.greater(k_expression, literal(v.val_str)) diff --git a/snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py b/snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py index 4ac58538a0..6fc281df92 100644 --- a/snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py +++ b/snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py @@ -21,8 +21,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 column, literal -from snuba.query.expressions import Expression +from snuba.query.dsl import column from snuba.query.logical import Query from snuba.query.query_settings import HTTPQuerySettings from snuba.reader import Row @@ -30,7 +29,11 @@ from snuba.web import QueryResult from snuba.web.query import run_query from snuba.web.rpc import RPCEndpoint -from snuba.web.rpc.common.common import base_conditions_and, treeify_or_and_conditions +from snuba.web.rpc.common.common import ( + base_conditions_and, + convert_filter_offset, + treeify_or_and_conditions, +) from snuba.web.rpc.common.debug_info import extract_response_meta from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException @@ -38,23 +41,6 @@ MAX_REQUEST_LIMIT = 1000 -def _convert_filter_offset(filter_offset: TraceItemFilter) -> Expression: - if not filter_offset.HasField("comparison_filter"): - raise TypeError("filter_offset needs to be a comparison filter") - if filter_offset.comparison_filter.op != ComparisonFilter.OP_GREATER_THAN: - raise TypeError("filter_offset must use the greater than comparison") - - k_expression = column(filter_offset.comparison_filter.key.name) - v = filter_offset.comparison_filter.value - value_type = v.WhichOneof("value") - if value_type != "val_str": - raise BadSnubaRPCRequestException( - "keys are strings, so please provide a string filter" - ) - - return f.greater(k_expression, literal(v.val_str)) - - def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaRequest: if req.limit > MAX_REQUEST_LIMIT: raise BadSnubaRPCRequestException( @@ -86,7 +72,7 @@ def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaReques base_conditions_and( req.meta, f.like(column("attr_key"), f"%{req.value_substring_match}%"), - _convert_filter_offset(req.page_token.filter_offset), + convert_filter_offset(req.page_token.filter_offset), ) if req.page_token.HasField("filter_offset") else base_conditions_and( diff --git a/snuba/web/rpc/v1/trace_item_attribute_values.py b/snuba/web/rpc/v1/trace_item_attribute_values.py index 3975b1acfc..e4bd5b20ae 100644 --- a/snuba/web/rpc/v1/trace_item_attribute_values.py +++ b/snuba/web/rpc/v1/trace_item_attribute_values.py @@ -7,6 +7,11 @@ TraceItemAttributeValuesResponse, ) from sentry_protos.snuba.v1.request_common_pb2 import PageToken +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey, AttributeValue +from sentry_protos.snuba.v1.trace_item_filter_pb2 import ( + ComparisonFilter, + TraceItemFilter, +) from snuba.attribution.appid import AppID from snuba.attribution.attribution_info import AttributionInfo @@ -17,6 +22,7 @@ from snuba.query.data_source.simple import Entity from snuba.query.dsl import Functions as f from snuba.query.dsl import column, literal, literals_array +from snuba.query.expressions import Expression from snuba.query.logical import Query from snuba.query.query_settings import HTTPQuerySettings from snuba.request import Request as SnubaRequest @@ -24,12 +30,52 @@ from snuba.web.rpc import RPCEndpoint from snuba.web.rpc.common.common import ( base_conditions_and, + convert_filter_offset, treeify_or_and_conditions, truncate_request_meta_to_day, ) from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException +def _build_base_conditions_and(request: TraceItemAttributeValuesRequest) -> Expression: + if request.value_substring_match is not None: + return ( + base_conditions_and( + request.meta, + f.equals(column("attr_key"), literal(request.key.name)), + # multiSearchAny has special treatment with ngram bloom filters + # https://clickhouse.com/docs/en/engines/table-engines/mergetree-family/mergetree#functions-support + f.multiSearchAny( + column("attr_value"), + literals_array(None, [literal(request.value_substring_match)]), + ), + convert_filter_offset(request.page_token.filter_offset), + ) + if request.page_token.HasField("filter_offset") + else base_conditions_and( + request.meta, + f.equals(column("attr_key"), literal(request.key.name)), + f.multiSearchAny( + column("attr_value"), + literals_array(None, [literal(request.value_substring_match)]), + ), + ) + ) + else: + return ( + base_conditions_and( + request.meta, + f.equals(column("attr_key"), literal(request.key.name)), + convert_filter_offset(request.page_token.filter_offset), + ) + if request.page_token.HasField("filter_offset") + else base_conditions_and( + request.meta, + f.equals(column("attr_key"), literal(request.key.name)), + ) + ) + + def _build_query(request: TraceItemAttributeValuesRequest) -> Query: if request.limit > 1000: raise BadSnubaRPCRequestException("Limit can be at most 1000") @@ -50,26 +96,8 @@ def _build_query(request: TraceItemAttributeValuesRequest) -> Query: expression=f.distinct(column("attr_value", alias="attr_value")), ), ], - condition=base_conditions_and( - request.meta, - f.equals(column("attr_key"), literal(request.key.name)), - # multiSearchAny has special treatment with ngram bloom filters - # https://clickhouse.com/docs/en/engines/table-engines/mergetree-family/mergetree#functions-support - f.multiSearchAny( - column("attr_value"), - literals_array(None, [literal(request.value_substring_match)]), - ), - ) - if request.value_substring_match is not None - else base_conditions_and( - request.meta, - f.equals(column("attr_key"), literal(request.key.name)), - ), + condition=_build_base_conditions_and(request), order_by=[ - OrderBy( - direction=OrderByDirection.ASC, expression=column("organization_id") - ), - OrderBy(direction=OrderByDirection.ASC, expression=column("attr_key")), OrderBy(direction=OrderByDirection.ASC, expression=column("attr_value")), ], limit=request.limit, @@ -115,12 +143,6 @@ def request_class(cls) -> Type[TraceItemAttributeValuesRequest]: def _execute( self, in_msg: TraceItemAttributeValuesRequest ) -> TraceItemAttributeValuesResponse: - if not in_msg.HasField("page_token"): - in_msg.page_token.offset = 0 - if in_msg.page_token.HasField("filter_offset"): - raise NotImplementedError( - "TraceItemAttributeValues does not currently support page_token.filter_offset, please use page_token.offset instead." - ) snuba_request = _build_snuba_request(in_msg) res = run_query( dataset=PluggableDataset(name="eap", all_entities=[]), @@ -130,5 +152,19 @@ def _execute( values = [r["attr_value"] for r in res.result.get("data", [])] return TraceItemAttributeValuesResponse( values=values, - page_token=PageToken(offset=in_msg.page_token.offset + len(values)), + page_token=( + PageToken(offset=in_msg.page_token.offset + len(values)) + if in_msg.page_token.HasField("offset") + else PageToken( + filter_offset=TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey( + type=AttributeKey.TYPE_STRING, name="attr_value" + ), + op=ComparisonFilter.OP_GREATER_THAN, + value=AttributeValue(val_str=values[-1]), + ) + ) + ) + ), ) diff --git a/tests/web/rpc/v1/test_trace_item_attribute_values_v1.py b/tests/web/rpc/v1/test_trace_item_attribute_values_v1.py index bd9d30b448..eeee20037d 100644 --- a/tests/web/rpc/v1/test_trace_item_attribute_values_v1.py +++ b/tests/web/rpc/v1/test_trace_item_attribute_values_v1.py @@ -7,9 +7,8 @@ from sentry_protos.snuba.v1.endpoint_trace_item_attributes_pb2 import ( TraceItemAttributeValuesRequest, ) -from sentry_protos.snuba.v1.request_common_pb2 import PageToken, RequestMeta +from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey -from sentry_protos.snuba.v1.trace_item_filter_pb2 import TraceItemFilter from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey @@ -147,15 +146,34 @@ def test_page_token(self, setup_teardown: Any) -> None: assert expected_values == [] def test_filter_offset_is_not_implemented_right_now(self) -> None: - req = TraceItemAttributeValuesRequest( - meta=COMMON_META, - limit=6, - key=AttributeKey(name="tag1", type=AttributeKey.TYPE_STRING), - value_substring_match="", - page_token=PageToken(filter_offset=TraceItemFilter()), - ) - with pytest.raises(NotImplementedError): - AttributeValuesRequest().execute(req) + expected_values = [ + "blah", + "derpderp", + "durp", + "herp", + "herpderp", + "some_last_value", + ] + total_number_of_values = len(expected_values) + + # grab 2 at a time until we get them all + done = 0 + page_token = None + at_a_time = 2 + while done < total_number_of_values: + req = TraceItemAttributeValuesRequest( + meta=COMMON_META, + limit=at_a_time, + key=AttributeKey(name="tag1", type=AttributeKey.TYPE_STRING), + value_substring_match="", + page_token=page_token, + ) + res = AttributeValuesRequest().execute(req) + page_token = res.page_token + assert res.values == expected_values[:at_a_time] + expected_values = expected_values[at_a_time:] + done += at_a_time + assert expected_values == [] def test_with_value_substring_match(self, setup_teardown: Any) -> None: message = TraceItemAttributeValuesRequest( From e8de78963fdba3491593c7bf8c76eaccf51cc65d Mon Sep 17 00:00:00 2001 From: Rachel Chen Date: Fri, 13 Dec 2024 09:53:47 -0800 Subject: [PATCH 2/2] test