feat(rpc): Update tags list rpc (#6301)

sentry-protos 0.1.21 changes the definition a bit, updating the endpoint
to accept it.

requirements.txt

@@ -45,4 +45,4 @@ sqlparse==0.4.2
 google-api-python-client==2.88.0
 sentry-usage-accountant==0.0.10
 freezegun==1.2.2
-sentry-protos==0.1.20
+sentry-protos==0.1.21

snuba/web/rpc/span_samples.py

@@ -72,7 +72,7 @@ def _build_query(request: SpanSamplesRequest) -> Query:
  limit=request.limit,
  )
  treeify_or_and_conditions(res)
- apply_virtual_columns(res, request.virtual_column_context)
+ apply_virtual_columns(res, request.virtual_column_contexts)
  return res
 
 

snuba/web/rpc/tags_list.py

@@ -2,9 +2,10 @@
 from typing import List, Optional
 
 from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
- TagsListRequest,
- TagsListResponse,
+ TraceItemAttributesRequest,
+ TraceItemAttributesResponse,
 )
+from sentry_protos.snuba.v1alpha.trace_item_attribute_pb2 import AttributeKey
 
 from snuba.clickhouse.formatter.nodes import FormattedQuery, StringNode
 from snuba.datasets.schemas.tables import TableSource
@@ -15,15 +16,17 @@
 
 
 def tags_list_query(
- request: TagsListRequest, _timer: Optional[Timer] = None
-) -> TagsListResponse:
- str_storage = get_storage(StorageKey("spans_str_attrs"))
- num_storage = get_storage(StorageKey("spans_num_attrs"))
-
- str_data_source = str_storage.get_schema().get_data_source()
- assert isinstance(str_data_source, TableSource)
- num_data_source = num_storage.get_schema().get_data_source()
- assert isinstance(num_data_source, TableSource)
+ request: TraceItemAttributesRequest, _timer: Optional[Timer] = None
+) -> TraceItemAttributesResponse:
+ if request.type == AttributeKey.Type.TYPE_STRING:
+ storage = get_storage(StorageKey("spans_str_attrs"))
+ elif request.type == AttributeKey.Type.TYPE_FLOAT:
+ storage = get_storage(StorageKey("spans_num_attrs"))
+ else:
+ return TraceItemAttributesResponse(tags=[])
+
+ data_source = storage.get_schema().get_data_source()
+ assert isinstance(data_source, TableSource)
 
  if request.limit > 1000:
  raise BadSnubaRPCRequestException("Limit can be at most 1000")
@@ -35,39 +38,26 @@ def tags_list_query(
  )
 
  query = f"""
-SELECT * FROM (
- SELECT DISTINCT attr_key, 'str' as type, timestamp
- FROM {str_data_source.get_table_name()}
- WHERE organization_id={request.meta.organization_id}
- AND project_id IN ({', '.join(str(pid) for pid in request.meta.project_ids)})
- AND timestamp BETWEEN fromUnixTimestamp({request.meta.start_timestamp.seconds}) AND fromUnixTimestamp({request.meta.end_timestamp.seconds})
-
- UNION ALL
-
- SELECT DISTINCT attr_key, 'num' as type, timestamp
- FROM {num_data_source.get_table_name()}
- WHERE organization_id={request.meta.organization_id}
- AND project_id IN ({', '.join(str(pid) for pid in request.meta.project_ids)})
- AND timestamp BETWEEN fromUnixTimestamp({request.meta.start_timestamp.seconds}) AND fromUnixTimestamp({request.meta.end_timestamp.seconds})
-)
+SELECT DISTINCT attr_key, timestamp
+FROM {data_source.get_table_name()}
+WHERE organization_id={request.meta.organization_id}
+AND project_id IN ({', '.join(str(pid) for pid in request.meta.project_ids)})
+AND timestamp BETWEEN fromUnixTimestamp({request.meta.start_timestamp.seconds}) AND fromUnixTimestamp({request.meta.end_timestamp.seconds})
 ORDER BY attr_key
 LIMIT {request.limit} OFFSET {request.offset}
 """
 
- cluster = str_storage.get_cluster()
+ cluster = storage.get_cluster()
  reader = cluster.get_reader()
  result = reader.execute(FormattedQuery([StringNode(query)]))
 
- tags: List[TagsListResponse.Tag] = []
+ tags: List[TraceItemAttributesResponse.Tag] = []
  for row in result.get("data", []):
  tags.append(
- TagsListResponse.Tag(
+ TraceItemAttributesResponse.Tag(
  name=row["attr_key"],
- type={
- "str": TagsListResponse.TYPE_STRING,
- "num": TagsListResponse.TYPE_NUMBER,
- }[row["type"]],
+ type=request.type,
  )
  )
 
- return TagsListResponse(tags=tags)
+ return TraceItemAttributesResponse(tags=tags)

snuba/web/views.py

@@ -40,7 +40,9 @@
  AggregateBucketRequest,
 )
 from sentry_protos.snuba.v1alpha.endpoint_span_samples_pb2 import SpanSamplesRequest
-from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import TagsListRequest
+from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
+ TraceItemAttributesRequest,
+)
 from werkzeug import Response as WerkzeugResponse
 from werkzeug.exceptions import InternalServerError
 
@@ -288,7 +290,7 @@ def rpc(*, name: str, timer: Timer) -> Response:
  ] = {
  "AggregateBucketRequest": (timeseries_query, AggregateBucketRequest),
  "SpanSamplesRequest": (span_samples_query, SpanSamplesRequest),
- "TagsListRequest": (tags_list_query, TagsListRequest),
+ "TraceItemAttributesRequest": (tags_list_query, TraceItemAttributesRequest),
  }
  try:
  endpoint, req_class = rpcs[name]

tests/web/rpc/test_span_samples.py

@@ -263,7 +263,7 @@ def test_with_virtual_columns(self, setup_teardown: Any) -> None:
  )
  ],
  limit=61,
- virtual_column_context=[
+ virtual_column_contexts=[
  VirtualColumnContext(
  from_column_name="project_id",
  to_column_name="project_name",
@@ -317,7 +317,7 @@ def test_order_by_virtual_columns(self, setup_teardown: Any) -> None:
  )
  ],
  limit=61,
- virtual_column_context=[
+ virtual_column_contexts=[
  VirtualColumnContext(
  from_column_name="color",
  to_column_name="special_color",

tests/web/rpc/test_tags_list.py

@@ -5,10 +5,11 @@
 import pytest
 from google.protobuf.timestamp_pb2 import Timestamp
 from sentry_protos.snuba.v1alpha.endpoint_tags_list_pb2 import (
- TagsListRequest,
- TagsListResponse,
+ TraceItemAttributesRequest,
+ TraceItemAttributesResponse,
 )
 from sentry_protos.snuba.v1alpha.request_common_pb2 import RequestMeta
+from sentry_protos.snuba.v1alpha.trace_item_attribute_pb2 import AttributeKey
 
 from snuba.datasets.storages.factory import get_storage
 from snuba.datasets.storages.storage_key import StorageKey
@@ -63,11 +64,11 @@ def setup_teardown(clickhouse_db: None, redis_db: None) -> None:
 
 @pytest.mark.clickhouse_db
 @pytest.mark.redis_db
-class TestTagsList(BaseApiTest):
+class TestTraceItemAttributes(BaseApiTest):
  def test_basic(self) -> None:
  ts = Timestamp()
  ts.GetCurrentTime()
- message = TagsListRequest(
+ message = TraceItemAttributesRequest(
  meta=RequestMeta(
  project_ids=[1, 2, 3],
  organization_id=1,
@@ -98,12 +99,12 @@ def test_basic(self) -> None:
  offset=20,
  )
  response = self.app.post(
- "/rpc/TagsListRequest", data=message.SerializeToString()
+ "/rpc/TraceItemAttributesRequest", data=message.SerializeToString()
  )
  assert response.status_code == 200
 
- def test_simple_case(self, setup_teardown: Any) -> None:
- message = TagsListRequest(
+ def test_simple_case_str(self, setup_teardown: Any) -> None:
+ message = TraceItemAttributesRequest(
  meta=RequestMeta(
  project_ids=[1, 2, 3],
  organization_id=1,
@@ -132,17 +133,58 @@ def test_simple_case(self, setup_teardown: Any) -> None:
  ),
  limit=10,
  offset=0,
+ type=AttributeKey.Type.TYPE_STRING,
  )
  response = tags_list_query(message)
  assert response.tags == [
- TagsListResponse.Tag(
- name=f"a_tag_{i:03}", type=TagsListResponse.TYPE_STRING
+ TraceItemAttributesResponse.Tag(
+ name=f"a_tag_{i:03}", type=AttributeKey.Type.TYPE_STRING
+ )
+ for i in range(0, 10)
+ ]
+
+ def test_simple_case_float(self, setup_teardown: Any) -> None:
+ message = TraceItemAttributesRequest(
+ 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 - 1,
+ 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()
+ )
+ ),
+ ),
+ limit=10,
+ offset=0,
+ type=AttributeKey.Type.TYPE_FLOAT,
+ )
+ response = tags_list_query(message)
+ assert response.tags == [
+ TraceItemAttributesResponse.Tag(
+ name=f"b_measurement_{i:03}", type=AttributeKey.Type.TYPE_FLOAT
  )
  for i in range(0, 10)
  ]
 
  def test_with_offset(self, setup_teardown: Any) -> None:
- message = TagsListRequest(
+ message = TraceItemAttributesRequest(
  meta=RequestMeta(
  project_ids=[1, 2, 3],
  organization_id=1,
@@ -170,21 +212,24 @@ def test_with_offset(self, setup_teardown: Any) -> None:
  ),
  ),
  limit=5,
- offset=29,
+ offset=10,
+ type=AttributeKey.Type.TYPE_FLOAT,
  )
  response = tags_list_query(message)
  assert response.tags == [
- TagsListResponse.Tag(name="a_tag_029", type=TagsListResponse.TYPE_STRING),
- TagsListResponse.Tag(
- name="b_measurement_000", type=TagsListResponse.TYPE_NUMBER
+ TraceItemAttributesResponse.Tag(
+ name="b_measurement_010", type=AttributeKey.Type.TYPE_FLOAT
+ ),
+ TraceItemAttributesResponse.Tag(
+ name="b_measurement_011", type=AttributeKey.Type.TYPE_FLOAT
  ),
- TagsListResponse.Tag(
- name="b_measurement_001", type=TagsListResponse.TYPE_NUMBER
+ TraceItemAttributesResponse.Tag(
+ name="b_measurement_012", type=AttributeKey.Type.TYPE_FLOAT
  ),
- TagsListResponse.Tag(
- name="b_measurement_002", type=TagsListResponse.TYPE_NUMBER
+ TraceItemAttributesResponse.Tag(
+ name="b_measurement_013", type=AttributeKey.Type.TYPE_FLOAT
  ),
- TagsListResponse.Tag(
- name="b_measurement_003", type=TagsListResponse.TYPE_NUMBER
+ TraceItemAttributesResponse.Tag(
+ name="b_measurement_014", type=AttributeKey.Type.TYPE_FLOAT
  ),
  ]