feat(eap): implement backend for "what tag values exist" RPC

snuba/web/rpc/__init__.py

@@ -0,0 +1,32 @@
+from typing import Any, Callable, Mapping, Tuple
+
+from google.protobuf.message import Message as ProtobufMessage
+from sentry_protos.snuba.v1alpha.endpoint_aggregate_bucket_pb2 import (
+    AggregateBucketRequest,
+)
+from sentry_protos.snuba.v1alpha.endpoint_span_samples_pb2 import SpanSamplesRequest
+from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
+    AttributeValuesRequest,
+    TraceItemAttributesRequest,
+)
+
+from snuba.utils.metrics.timer import Timer
+from snuba.web.rpc.span_samples import span_samples_query
+from snuba.web.rpc.timeseries import timeseries_query
+from snuba.web.rpc.trace_item_attribute_list import trace_item_attribute_list_query
+from snuba.web.rpc.trace_item_attribute_values import trace_item_attribute_values_query
+
+ALL_RPCS: Mapping[
+    str, Tuple[Callable[[Any, Timer], ProtobufMessage], type[ProtobufMessage]]
+] = {
+    "AggregateBucketRequest": (timeseries_query, AggregateBucketRequest),
+    "SpanSamplesRequest": (span_samples_query, SpanSamplesRequest),
+    "TraceItemAttributesRequest": (
+        trace_item_attribute_list_query,
+        TraceItemAttributesRequest,
+    ),
+    "AttributeValuesRequest": (
+        trace_item_attribute_values_query,
+        AttributeValuesRequest,
+    ),
+}

snuba/web/rpc/tags_list.py → snuba/web/rpc/trace_item_attribute_list.py

@@ -15,7 +15,7 @@
 from snuba.web.rpc.exceptions import BadSnubaRPCRequestException
 
 
-def tags_list_query(
+def trace_item_attribute_list_query(
     request: TraceItemAttributesRequest, _timer: Optional[Timer] = None
 ) -> TraceItemAttributesResponse:
     if request.type == AttributeKey.Type.TYPE_STRING:
@@ -31,11 +31,16 @@ def tags_list_query(
     if request.limit > 1000:
         raise BadSnubaRPCRequestException("Limit can be at most 1000")
 
+    # this table stores timestamp as toStartOfDay(x) in UTC, so if you request 4PM - 8PM on a specific day, nada
     start_timestamp = datetime.utcfromtimestamp(request.meta.start_timestamp.seconds)
-    if start_timestamp.day >= datetime.utcnow().day and start_timestamp.hour != 0:
-        raise BadSnubaRPCRequestException(
-            "Tags' timestamps are stored per-day, you probably want to set start_timestamp to UTC 00:00 today or a time yesterday."
+    end_timestamp = datetime.utcfromtimestamp(request.meta.end_timestamp.seconds)
+    if start_timestamp.day == end_timestamp.day:
+        start_timestamp = start_timestamp.replace(
+            day=start_timestamp.day - 1, hour=0, minute=0, second=0, microsecond=0
         )
+        end_timestamp = end_timestamp.replace(day=end_timestamp.day + 1)
+        request.meta.start_timestamp.seconds = int(start_timestamp.timestamp())
+        request.meta.end_timestamp.seconds = int(end_timestamp.timestamp())
 
     query = f"""
 SELECT DISTINCT attr_key, timestamp

snuba/web/rpc/trace_item_attribute_values.py

@@ -0,0 +1,115 @@
+import uuid
+from datetime import datetime
+
+from google.protobuf.json_format import MessageToDict
+from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
+    AttributeValuesRequest,
+    AttributeValuesResponse,
+)
+
+from snuba.attribution.appid import AppID
+from snuba.attribution.attribution_info import AttributionInfo
+from snuba.datasets.entities.entity_key import EntityKey
+from snuba.datasets.entities.factory import get_entity
+from snuba.datasets.pluggable_dataset import PluggableDataset
+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, literals_array
+from snuba.query.logical import Query
+from snuba.query.query_settings import HTTPQuerySettings
+from snuba.request import Request as SnubaRequest
+from snuba.utils.metrics.timer import Timer
+from snuba.web.query import run_query
+from snuba.web.rpc.common import base_conditions_and, treeify_or_and_conditions
+from snuba.web.rpc.exceptions import BadSnubaRPCRequestException
+
+
+def _build_query(request: AttributeValuesRequest) -> Query:
+    if request.limit > 1000:
+        raise BadSnubaRPCRequestException("Limit can be at most 1000")
+
+    entity = Entity(
+        key=EntityKey("spans_str_attrs"),
+        schema=get_entity(EntityKey("spans_str_attrs")).get_data_model(),
+        sample=None,
+    )
+
+    # this table stores timestamp as toStartOfDay(x) in UTC, so if you request 4PM - 8PM on a specific day, nada
+    start_timestamp = datetime.utcfromtimestamp(request.meta.start_timestamp.seconds)
+    end_timestamp = datetime.utcfromtimestamp(request.meta.end_timestamp.seconds)
+    if start_timestamp.day == end_timestamp.day:
+        start_timestamp = start_timestamp.replace(
+            day=start_timestamp.day - 1, hour=0, minute=0, second=0, microsecond=0
+        )
+        end_timestamp = end_timestamp.replace(day=end_timestamp.day + 1)
+        request.meta.start_timestamp.seconds = int(start_timestamp.timestamp())
+        request.meta.end_timestamp.seconds = int(end_timestamp.timestamp())
+
+    res = Query(
+        from_clause=entity,
+        selected_columns=[
+            SelectedExpression(
+                name="attr_value",
+                expression=f.distinct(column("attr_value", alias="attr_value")),
+            ),
+        ],
+        condition=base_conditions_and(
+            request.meta,
+            f.equals(column("attr_key"), literal(request.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)]),
+            ),
+        ),
+        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,
+        offset=request.offset,
+    )
+    treeify_or_and_conditions(res)
+    return res
+
+
+def _build_snuba_request(
+    request: AttributeValuesRequest,
+) -> SnubaRequest:
+    return SnubaRequest(
+        id=str(uuid.uuid4()),
+        original_body=MessageToDict(request),
+        query=_build_query(request),
+        query_settings=HTTPQuerySettings(),
+        attribution_info=AttributionInfo(
+            referrer=request.meta.referrer,
+            team="eap",
+            feature="eap",
+            tenant_ids={
+                "organization_id": request.meta.organization_id,
+                "referrer": request.meta.referrer,
+            },
+            app_id=AppID("eap"),
+            parent_api="trace_item_values",
+        ),
+    )
+
+
+def trace_item_attribute_values_query(
+    request: AttributeValuesRequest, timer: Timer | None = None
+) -> AttributeValuesResponse:
+    timer = timer or Timer("trace_item_values")
+    snuba_request = _build_snuba_request(request)
+    res = run_query(
+        dataset=PluggableDataset(name="eap", all_entities=[]),
+        request=snuba_request,
+        timer=timer,
+    )
+    return AttributeValuesResponse(
+        values=[r["attr_value"] for r in res.result.get("data", [])]
+    )

snuba/web/views.py

@@ -35,14 +35,6 @@
     render_template,
 )
 from flask import request as http_request
-from google.protobuf.message import Message as ProtobufMessage
-from sentry_protos.snuba.v1alpha.endpoint_aggregate_bucket_pb2 import (
-    AggregateBucketRequest,
-)
-from sentry_protos.snuba.v1alpha.endpoint_span_samples_pb2 import SpanSamplesRequest
-from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
-    TraceItemAttributesRequest,
-)
 from werkzeug import Response as WerkzeugResponse
 from werkzeug.exceptions import InternalServerError
 
@@ -83,10 +75,8 @@
 from snuba.web.converters import DatasetConverter, EntityConverter, StorageConverter
 from snuba.web.delete_query import DeletesNotEnabledError, delete_from_storage
 from snuba.web.query import parse_and_run_query
+from snuba.web.rpc import ALL_RPCS
 from snuba.web.rpc.exceptions import BadSnubaRPCRequestException
-from snuba.web.rpc.span_samples import span_samples_query as span_samples_query
-from snuba.web.rpc.tags_list import tags_list_query
-from snuba.web.rpc.timeseries import timeseries_query as timeseries_query
 from snuba.writer import BatchWriterEncoderWrapper, WriterTableRow
 
 logger = logging.getLogger("snuba.api")
@@ -285,15 +275,8 @@ def unqualified_query_view(*, timer: Timer) -> Union[Response, str, WerkzeugResp
 @application.route("/rpc/<name>", methods=["POST"])
 @util.time_request("timeseries")
 def rpc(*, name: str, timer: Timer) -> Response:
-    rpcs: Mapping[
-        str, Tuple[Callable[[Any, Timer], ProtobufMessage], type[ProtobufMessage]]
-    ] = {
-        "AggregateBucketRequest": (timeseries_query, AggregateBucketRequest),
-        "SpanSamplesRequest": (span_samples_query, SpanSamplesRequest),
-        "TraceItemAttributesRequest": (tags_list_query, TraceItemAttributesRequest),
-    }
     try:
-        endpoint, req_class = rpcs[name]
+        endpoint, req_class = ALL_RPCS[name]
 
         req = req_class()
         req.ParseFromString(http_request.data)

tests/web/rpc/test_tags_list.py → tests/web/rpc/test_trace_item_attribute_list.py

@@ -13,7 +13,7 @@
 
 from snuba.datasets.storages.factory import get_storage
 from snuba.datasets.storages.storage_key import StorageKey
-from snuba.web.rpc.tags_list import tags_list_query
+from snuba.web.rpc.trace_item_attribute_list import trace_item_attribute_list_query
 from tests.base import BaseApiTest
 from tests.helpers import write_raw_unprocessed_events
 
@@ -135,7 +135,7 @@ def test_simple_case_str(self, setup_teardown: Any) -> None:
             offset=0,
             type=AttributeKey.Type.TYPE_STRING,
         )
-        response = tags_list_query(message)
+        response = trace_item_attribute_list_query(message)
         assert response.tags == [
             TraceItemAttributesResponse.Tag(
                 name=f"a_tag_{i:03}", type=AttributeKey.Type.TYPE_STRING
@@ -175,7 +175,7 @@ def test_simple_case_float(self, setup_teardown: Any) -> None:
             offset=0,
             type=AttributeKey.Type.TYPE_FLOAT,
         )
-        response = tags_list_query(message)
+        response = trace_item_attribute_list_query(message)
         assert response.tags == [
             TraceItemAttributesResponse.Tag(
                 name=f"b_measurement_{i:03}", type=AttributeKey.Type.TYPE_FLOAT
@@ -215,7 +215,7 @@ def test_with_offset(self, setup_teardown: Any) -> None:
             offset=10,
             type=AttributeKey.Type.TYPE_FLOAT,
         )
-        response = tags_list_query(message)
+        response = trace_item_attribute_list_query(message)
         assert response.tags == [
             TraceItemAttributesResponse.Tag(
                 name="b_measurement_010", type=AttributeKey.Type.TYPE_FLOAT

tests/web/rpc/test_trace_item_attribute_values.py

@@ -0,0 +1,127 @@
+import uuid
+from datetime import UTC, datetime, timedelta
+from typing import Any, Mapping
+
+import pytest
+from google.protobuf.timestamp_pb2 import Timestamp
+from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import AttributeValuesRequest
+from sentry_protos.snuba.v1alpha.request_common_pb2 import RequestMeta
+
+from snuba.datasets.storages.factory import get_storage
+from snuba.datasets.storages.storage_key import StorageKey
+from snuba.web.rpc.trace_item_attribute_values import trace_item_attribute_values_query
+from tests.base import BaseApiTest
+from tests.helpers import write_raw_unprocessed_events
+
+BASE_TIME = datetime.utcnow().replace(minute=0, second=0, microsecond=0) - timedelta(
+    minutes=180
+)
+COMMON_META = RequestMeta(
+    project_ids=[1, 2, 3],
+    organization_id=1,
+    cogs_category="something",
+    referrer="something",
+    start_timestamp=Timestamp(
+        seconds=int(
+            datetime(
+                year=BASE_TIME.year,
+                month=BASE_TIME.month,
+                day=BASE_TIME.day,
+                tzinfo=UTC,
+            ).timestamp()
+        )
+    ),
+    end_timestamp=Timestamp(
+        seconds=int(
+            datetime(
+                year=BASE_TIME.year,
+                month=BASE_TIME.month,
+                day=BASE_TIME.day + 1,
+                tzinfo=UTC,
+            ).timestamp()
+        )
+    ),
+)
+
+
+def gen_message(tags: Mapping[str, str]) -> Mapping[str, Any]:
+    return {
+        "description": "/api/0/relays/projectconfigs/",
+        "duration_ms": 152,
+        "event_id": "d826225de75d42d6b2f01b957d51f18f",
+        "exclusive_time_ms": 0.228,
+        "is_segment": True,
+        "data": {},
+        "measurements": {},
+        "organization_id": 1,
+        "origin": "auto.http.django",
+        "project_id": 1,
+        "received": 1721319572.877828,
+        "retention_days": 90,
+        "segment_id": "8873a98879faf06d",
+        "sentry_tags": {},
+        "span_id": uuid.uuid4().hex,
+        "tags": tags,
+        "trace_id": uuid.uuid4().hex,
+        "start_timestamp_ms": int(BASE_TIME.timestamp() * 1000),
+        "start_timestamp_precise": BASE_TIME.timestamp(),
+        "end_timestamp_precise": BASE_TIME.timestamp() + 1,
+    }
+
+
+@pytest.fixture(autouse=True)
+def setup_teardown(clickhouse_db: None, redis_db: None) -> None:
+    spans_storage = get_storage(StorageKey("eap_spans"))
+    messages = [
+        gen_message({"tag1": "herp", "tag2": "herp"}),
+        gen_message({"tag1": "herpderp", "tag2": "herp"}),
+        gen_message({"tag1": "durp", "tag3": "herp"}),
+        gen_message({"tag1": "blah", "tag2": "herp"}),
+        gen_message({"tag1": "derpderp", "tag2": "derp"}),
+        gen_message({"tag2": "hehe"}),
+        gen_message({"tag1": "some_last_value"}),
+    ]
+    write_raw_unprocessed_events(spans_storage, messages)  # type: ignore
+
+
+@pytest.mark.clickhouse_db
+@pytest.mark.redis_db
+class TestTraceItemAttributes(BaseApiTest):
+    def test_basic(self) -> None:
+        ts = Timestamp()
+        ts.GetCurrentTime()
+        message = AttributeValuesRequest(
+            meta=COMMON_META,
+            name="tag1",
+            limit=10,
+            offset=20,
+        )
+        response = self.app.post(
+            "/rpc/AttributeValuesRequest", data=message.SerializeToString()
+        )
+        assert response.status_code == 200
+
+    def test_simple_case(self, setup_teardown: Any) -> None:
+        message = AttributeValuesRequest(meta=COMMON_META, limit=5, name="tag1")
+        response = trace_item_attribute_values_query(message)
+        assert response.values == ["blah", "derpderp", "durp", "herp", "herpderp"]
+
+    def test_offset(self, setup_teardown: Any) -> None:
+        message = AttributeValuesRequest(
+            meta=COMMON_META, limit=5, offset=1, name="tag1"
+        )
+        response = trace_item_attribute_values_query(message)
+        assert response.values == [
+            "derpderp",
+            "durp",
+            "herp",
+            "herpderp",
+            "some_last_value",
+        ]
+
+    def test_with_value_filter(self, setup_teardown: Any) -> None:
+        message = AttributeValuesRequest(
+            meta=COMMON_META, limit=5, name="tag1", value_substring_match="erp"
+        )
+        response = trace_item_attribute_values_query(message)
+        assert response.values == ["derpderp", "herp", "herpderp"]