Implement filter offset for attribute values API

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))

snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py

@@ -21,40 +21,26 @@
 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
 from snuba.request import Request as SnubaRequest
 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
 
 # max value the user can provide for 'limit' in their request
 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(

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,19 +22,60 @@
 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
 from snuba.web.query import run_query
 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]),
+                        )
+                    )
+                )
+            ),
         )

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(