feat(attribute-values): Add counts to response and orderby

snuba/web/rpc/v1/trace_item_attribute_values.py

@@ -7,11 +7,7 @@
     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 sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
 
 from snuba.attribution.appid import AppID
 from snuba.attribution.attribution_info import AttributionInfo
@@ -25,7 +21,7 @@
 from snuba.query.data_source.simple import Entity, LogicalDataSource
 from snuba.query.dsl import Functions as f
 from snuba.query.dsl import column
-from snuba.query.expressions import Expression
+from snuba.query.expressions import Expression, FunctionCall
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
@@ -126,12 +122,22 @@ def _build_query(
         selected_columns=[
             SelectedExpression(
                 name="attr_value",
-                expression=f.distinct(column(attr_value.alias, alias="attr_value")),
+                expression=column(attr_value.alias, alias="attr_value"),
+            ),
+            SelectedExpression(
+                name="count()",
+                expression=FunctionCall(
+                    alias="count()",
+                    function_name="count",
+                    parameters=(),
+                ),
             ),
         ],
         order_by=[
+            OrderBy(direction=OrderByDirection.DESC, expression=column("count()")),
             OrderBy(direction=OrderByDirection.ASC, expression=column("attr_value")),
         ],
+        groupby=[column("attr_value")],
         limit=request.limit,
         offset=(request.page_token.offset if request.page_token.HasField("offset") else 0),
     )
@@ -184,6 +190,7 @@ def _execute(self, in_msg: TraceItemAttributeValuesRequest) -> TraceItemAttribut
         if in_msg.key.name == "sentry.item_id" and in_msg.value_substring_match:
             return TraceItemAttributeValuesResponse(
                 values=[in_msg.value_substring_match],
+                counts=[1],
                 page_token=None,
             )
         in_msg.limit = in_msg.limit or 1000
@@ -193,25 +200,23 @@ def _execute(self, in_msg: TraceItemAttributeValuesRequest) -> TraceItemAttribut
             request=snuba_request,
             timer=self._timer,
         )
-        values = [r["attr_value"] for r in res.result.get("data", [])]
+        values, counts = [], []
+        for row in res.result.get("data", []):
+            values.append(row["attr_value"])
+            counts.append(row.get("count()", 0))
         if len(values) == 0:
             return TraceItemAttributeValuesResponse(
                 values=values,
+                counts=counts,
                 page_token=None,
             )
         return TraceItemAttributeValuesResponse(
             values=values,
+            counts=counts,
             page_token=(
-                PageToken(offset=in_msg.page_token.offset + len(values))
-                if in_msg.page_token.HasField("offset") or len(values) == 0
-                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]),
-                        )
-                    )
+                PageToken(
+                    offset=(in_msg.page_token.offset if in_msg.page_token.HasField("offset") else 0)
+                    + len(values)
                 )
             ),
         )

tests/web/rpc/v1/test_trace_item_attribute_values_v1.py

@@ -93,6 +93,13 @@ def setup_teardown(eap: None, redis_db: None) -> Generator[List[bytes], None, No
                 "tag2": AnyValue(string_value="derp"),
             },
         ),
+        gen_item_message(
+            start_timestamp=start_timestamp,
+            attributes={
+                "tag1": AnyValue(string_value="derpderp"),
+                "tag2": AnyValue(string_value="derp"),
+            },
+        ),
         gen_item_message(
             start_timestamp=start_timestamp,
             attributes={"tag2": AnyValue(string_value="hehe")},
@@ -135,7 +142,8 @@ def test_simple_case(self, setup_teardown: Any) -> None:
             key=AttributeKey(name="tag1", type=AttributeKey.TYPE_STRING),
         )
         response = AttributeValuesRequest().execute(message)
-        assert response.values == ["blah", "derpderp", "durp", "herp", "herpderp"]
+        assert response.values == ["derpderp", "blah", "durp", "herp", "herpderp"]
+        assert response.counts == [2, 1, 1, 1, 1]
 
     def test_with_value_substring_match(self, setup_teardown: Any) -> None:
         message = TraceItemAttributeValuesRequest(
@@ -146,6 +154,7 @@ def test_with_value_substring_match(self, setup_teardown: Any) -> None:
         )
         response = AttributeValuesRequest().execute(message)
         assert response.values == ["derpderp", "herp", "herpderp"]
+        assert response.counts == [2, 1, 1]
 
     def test_empty_results(self) -> None:
         req = TraceItemAttributeValuesRequest(
@@ -162,6 +171,7 @@ def test_empty_results(self) -> None:
         )
         res = AttributeValuesRequest().execute(req)
         assert res.values == []
+        assert res.counts == []
 
     def test_item_id_substring_match(self, setup_teardown: List[bytes]) -> None:
         first_msg_bytes = setup_teardown[0]
@@ -182,6 +192,7 @@ def test_item_id_substring_match(self, setup_teardown: List[bytes]) -> None:
         )
         res = AttributeValuesRequest().execute(req)
         assert res.values == [item_id]
+        assert res.counts == [1]
 
     def test_deprecated_alias_attribute(self) -> None:
         """db.system.name request returns values stored only under deprecated key db.system."""
@@ -218,3 +229,25 @@ def test_deprecated_alias_attribute(self) -> None:
         )
         response = AttributeValuesRequest().execute(message)
         assert sorted(response.values) == ["postgresql", "redis"]
+        assert response.counts == [1, 1]
+
+    def test_pagination(self, setup_teardown: Any) -> None:
+        message = TraceItemAttributeValuesRequest(
+            meta=COMMON_META,
+            limit=1,
+            key=AttributeKey(name="tag1", type=AttributeKey.TYPE_STRING),
+        )
+        response = AttributeValuesRequest().execute(message)
+        assert response.values == ["derpderp"]
+        assert response.counts == [2]
+
+        for expected in ["blah", "durp", "herp", "herpderp"]:
+            message = TraceItemAttributeValuesRequest(
+                meta=COMMON_META,
+                limit=1,
+                key=AttributeKey(name="tag1", type=AttributeKey.TYPE_STRING),
+                page_token=response.page_token,
+            )
+            response = AttributeValuesRequest().execute(message)
+            assert response.values == [expected]
+            assert response.counts == [1]