Implement filter offset for attribute keys API (#6618)

https://github.com/getsentry/projects/issues/376 Need to do this for attribute values API next --------- Co-authored-by: Rachel Chen <[email protected]>
getsentry · Dec 11, 2024 · 353c898 · 353c898
1 parent 6960a57
commit 353c898
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 34 deletions.
diff --git a/snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py b/snuba/web/rpc/v1/endpoint_trace_item_attribute_names.py
@@ -7,7 +7,11 @@
     TraceItemAttributeNamesResponse,
 )
 from sentry_protos.snuba.v1.request_common_pb2 import PageToken
-from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
+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,7 +21,8 @@
 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
+from snuba.query.dsl import column, literal
+from snuba.query.expressions import Expression
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.reader import Row
@@ -33,6 +38,23 @@
 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(
@@ -60,6 +82,18 @@ def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaReques
             f"Attribute type '{req.type}' is not supported. Supported types are: TYPE_STRING, TYPE_FLOAT, TYPE_INT, TYPE_BOOLEAN"
         )
 
+    condition = (
+        base_conditions_and(
+            req.meta,
+            f.like(column("attr_key"), f"%{req.value_substring_match}%"),
+            _convert_filter_offset(req.page_token.filter_offset),
+        )
+        if req.page_token.HasField("filter_offset")
+        else base_conditions_and(
+            req.meta, f.like(column("attr_key"), f"%{req.value_substring_match}%")
+        )
+    )
+
     query = Query(
         from_clause=entity,
         selected_columns=[
@@ -68,9 +102,7 @@ def convert_to_snuba_request(req: TraceItemAttributeNamesRequest) -> SnubaReques
                 expression=column("attr_key"),
             ),
         ],
-        condition=base_conditions_and(
-            req.meta, f.like(column("attr_key"), f"%{req.value_substring_match}%")
-        ),
+        condition=condition,
         order_by=[
             OrderBy(direction=OrderByDirection.ASC, expression=column("attr_key")),
         ],
@@ -138,9 +170,24 @@ def _build_response(
         res: QueryResult,
     ) -> TraceItemAttributeNamesResponse:
         attributes = convert_to_attributes(res, req.type)
+        page_token = (
+            PageToken(offset=req.page_token.offset + len(attributes))
+            if req.page_token.HasField("offset")
+            else PageToken(
+                filter_offset=TraceItemFilter(
+                    comparison_filter=ComparisonFilter(
+                        key=AttributeKey(
+                            type=AttributeKey.TYPE_STRING, name="attr_key"
+                        ),
+                        op=ComparisonFilter.OP_GREATER_THAN,
+                        value=AttributeValue(val_str=attributes[-1].name),
+                    )
+                )
+            )
+        )
         return TraceItemAttributeNamesResponse(
             attributes=attributes,
-            page_token=PageToken(offset=req.page_token.offset + len(attributes)),
+            page_token=page_token,
             meta=extract_response_meta(
                 req.meta.request_id, req.meta.debug, [res], [self._timer]
             ),
@@ -149,12 +196,6 @@ def _build_response(
     def _execute(
         self, req: TraceItemAttributeNamesRequest
     ) -> TraceItemAttributeNamesResponse:
-        if not req.HasField("page_token"):
-            req.page_token.offset = 0
-        if req.page_token.HasField("filter_offset"):
-            raise NotImplementedError(
-                "EndpointTraceItemAttributeNames does not currently support page_token.filter_offset, please use page_token.offset instead."
-            )
         if not req.meta.request_id:
             req.meta.request_id = str(uuid.uuid4())
 

diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_attribute_names.py b/tests/web/rpc/v1/test_endpoint_trace_item_attribute_names.py
@@ -10,7 +10,6 @@
 )
 from sentry_protos.snuba.v1.request_common_pb2 import PageToken, 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
@@ -171,7 +170,7 @@ def test_with_filter(self) -> None:
         ]
         assert res.attributes == expected
 
-    def test_with_page_token(self) -> None:
+    def test_with_page_token_offset(self) -> None:
         # this is all the expected attributes
         expected_attributes = []
         for i in range(TOTAL_GENERATED_ATTR_PER_TYPE):
@@ -183,7 +182,7 @@ def test_with_page_token(self) -> None:
             )
         # grab 10 at a time until we get them all
         done = 0
-        page_token = None
+        page_token = PageToken(offset=0)
         at_a_time = 10
         while done < TOTAL_GENERATED_ATTR_PER_TYPE:
             req = TraceItemAttributeNamesRequest(
@@ -212,26 +211,45 @@ def test_with_page_token(self) -> None:
         assert expected_attributes == []
 
     def test_page_token_offset_filter(self) -> None:
-        req = TraceItemAttributeNamesRequest(
-            meta=RequestMeta(
-                project_ids=[1, 2, 3],
-                organization_id=1,
-                cogs_category="something",
-                referrer="something",
-                start_timestamp=Timestamp(
-                    seconds=int((BASE_TIME - timedelta(days=1)).timestamp())
-                ),
-                end_timestamp=Timestamp(
-                    seconds=int((BASE_TIME + timedelta(days=1)).timestamp())
+
+        expected_attributes = []
+        for i in range(TOTAL_GENERATED_ATTR_PER_TYPE):
+            expected_attributes.append(
+                TraceItemAttributeNamesResponse.Attribute(
+                    name=f"a_tag_{str(i).zfill(3)}",
+                    type=AttributeKey.Type.TYPE_STRING,
+                )
+            )
+        # grab 10 at a time until we get them all
+        done = 0
+        page_token = None
+        at_a_time = 10
+
+        while done < TOTAL_GENERATED_ATTR_PER_TYPE:
+            req = TraceItemAttributeNamesRequest(
+                meta=RequestMeta(
+                    project_ids=[1, 2, 3],
+                    organization_id=1,
+                    cogs_category="something",
+                    referrer="something",
+                    start_timestamp=Timestamp(
+                        seconds=int((BASE_TIME - timedelta(days=1)).timestamp())
+                    ),
+                    end_timestamp=Timestamp(
+                        seconds=int((BASE_TIME + timedelta(days=1)).timestamp())
+                    ),
                 ),
-            ),
-            limit=10,
-            type=AttributeKey.Type.TYPE_STRING,
-            value_substring_match="",
-            page_token=PageToken(filter_offset=TraceItemFilter()),
-        )
-        with pytest.raises(NotImplementedError):
-            EndpointTraceItemAttributeNames().execute(req)
+                limit=at_a_time,
+                type=AttributeKey.Type.TYPE_STRING,
+                value_substring_match="",
+                page_token=page_token,
+            )
+            res = EndpointTraceItemAttributeNames().execute(req)
+            page_token = res.page_token
+            assert res.attributes == expected_attributes[:at_a_time]
+            expected_attributes = expected_attributes[at_a_time:]
+            done += at_a_time
+        assert expected_attributes == []
 
     def test_response_metadata(self) -> None:
         # debug must be true in RequestMeta for it to return query_info in the response