diff --git a/snuba/protos/common.py b/snuba/protos/common.py index b764901576..9fdf5e7a08 100644 --- a/snuba/protos/common.py +++ b/snuba/protos/common.py @@ -1,6 +1,6 @@ from collections import defaultdict from collections.abc import Mapping, Sequence -from typing import Final +from typing import Final, NamedTuple from sentry_conventions.attributes import ATTRIBUTE_METADATA from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey @@ -26,23 +26,42 @@ class MalformedAttributeException(Exception): pass -COLUMN_PREFIX: str = "sentry." - -NORMALIZED_COLUMNS_EAP_ITEMS: Final[Mapping[str, Sequence[AttributeKey.Type.ValueType]]] = { - f"{COLUMN_PREFIX}organization_id": [AttributeKey.Type.TYPE_INT], - f"{COLUMN_PREFIX}project_id": [AttributeKey.Type.TYPE_INT], - f"{COLUMN_PREFIX}timestamp": [ - AttributeKey.Type.TYPE_FLOAT, - AttributeKey.Type.TYPE_DOUBLE, - AttributeKey.Type.TYPE_INT, - AttributeKey.Type.TYPE_STRING, - ], - f"{COLUMN_PREFIX}trace_id": [ - AttributeKey.Type.TYPE_STRING - ], # this gets converted from a uuid to a string in a storage processor - f"{COLUMN_PREFIX}item_id": [AttributeKey.Type.TYPE_STRING], - f"{COLUMN_PREFIX}sampling_weight": [AttributeKey.Type.TYPE_DOUBLE], - f"{COLUMN_PREFIX}sampling_factor": [AttributeKey.Type.TYPE_DOUBLE], +class NormalizedColumn(NamedTuple): + column_name: str + types: Sequence[AttributeKey.Type.ValueType] + + +def sentry_column(column_name: str) -> str: + return f"sentry.{column_name}" + + +NORMALIZED_COLUMNS_EAP_ITEMS: Final[Mapping[str, NormalizedColumn]] = { + sentry_column("organization_id"): NormalizedColumn( + "organization_id", [AttributeKey.Type.TYPE_INT] + ), + sentry_column("project_id"): NormalizedColumn("project_id", [AttributeKey.Type.TYPE_INT]), + sentry_column("timestamp"): NormalizedColumn( + "timestamp", + [ + AttributeKey.Type.TYPE_FLOAT, + AttributeKey.Type.TYPE_DOUBLE, + AttributeKey.Type.TYPE_INT, + AttributeKey.Type.TYPE_STRING, + ], + ), + # trace_id gets converted from a uuid to a string in a storage processor + sentry_column("trace_id"): NormalizedColumn("trace_id", [AttributeKey.Type.TYPE_STRING]), + sentry_column("item_id"): NormalizedColumn("item_id", [AttributeKey.Type.TYPE_STRING]), + sentry_column("sampling_weight"): NormalizedColumn( + "sampling_weight", [AttributeKey.Type.TYPE_DOUBLE] + ), + sentry_column("sampling_factor"): NormalizedColumn( + "sampling_factor", [AttributeKey.Type.TYPE_DOUBLE] + ), + sentry_column("session_id"): NormalizedColumn("session_id", [AttributeKey.Type.TYPE_STRING]), + "gen_ai.conversation.id": NormalizedColumn( + "ai_conversation_id", [AttributeKey.Type.TYPE_STRING] + ), } PROTO_TYPE_TO_CLICKHOUSE_TYPE: Final[Mapping[AttributeKey.Type.ValueType, str]] = { @@ -172,6 +191,12 @@ def _generate_subscriptable_reference( "attributes_array_bool", ) +# Normalized String columns whose unset value ingests as an empty string (not NULL) — e.g. +# ai_conversation_id when a TraceItem has no conversation_id. Because the column is never NULL, +# existence must be checked with notEmpty(...) instead of isNotNull(...), which would otherwise +# match every row (see get_field_existence_expression). +EMPTY_STRING_DEFAULT_COLUMNS: tuple[str, ...] = ("ai_conversation_id",) + def type_array_typed_columns_select_expressions(attr_key: AttributeKey) -> list[FunctionCall]: """Native ``arrayElement`` read per typed array map column, for a deprecated untyped @@ -297,11 +322,12 @@ def attribute_key_to_expression(attr_key: AttributeKey) -> Expression: return column("attr_key") if attr_key.name in NORMALIZED_COLUMNS_EAP_ITEMS: - if attr_key.type not in NORMALIZED_COLUMNS_EAP_ITEMS[attr_key.name]: + normalized_column = NORMALIZED_COLUMNS_EAP_ITEMS[attr_key.name] + if attr_key.type not in normalized_column.types: formatted_attribute_types = ", ".join( map( AttributeKey.Type.Name, - NORMALIZED_COLUMNS_EAP_ITEMS[attr_key.name], + normalized_column.types, ) ) raise MalformedAttributeException( @@ -309,7 +335,7 @@ def attribute_key_to_expression(attr_key: AttributeKey) -> Expression: ) return f.cast( - column(attr_key.name[len(COLUMN_PREFIX) :]), + column(normalized_column.column_name), PROTO_TYPE_TO_CLICKHOUSE_TYPE[attr_key.type], alias=alias, ) diff --git a/snuba/web/rpc/common/common.py b/snuba/web/rpc/common/common.py index b0fbc42887..50b7b46eb7 100644 --- a/snuba/web/rpc/common/common.py +++ b/snuba/web/rpc/common/common.py @@ -17,13 +17,14 @@ from snuba.protos.common import ( ARRAY_TYPES, ATTRIBUTES_TO_COALESCE, - COLUMN_PREFIX, + EMPTY_STRING_DEFAULT_COLUMNS, NORMALIZED_COLUMNS_EAP_ITEMS, PROTO_TYPE_TO_ATTRIBUTE_COLUMN, PROTO_TYPE_TO_CLICKHOUSE_TYPE, TYPED_ARRAY_MAP_COLUMNS, MalformedAttributeException, array_element_column, + sentry_column, type_array_to_membership_array_expression_from_typed_columns, type_array_typed_column_native_array, ) @@ -922,7 +923,7 @@ def trace_item_filters_to_expression( # index- and partition-prunable. We reuse timestamp_seconds_to_datetime_literal # so a bound equal to the mandatory range is byte-identical to it and gets # collapsed by dedupe_timestamp_conditions. - if k.name == f"{COLUMN_PREFIX}timestamp" and value_type in ( + if k.name == sentry_column("timestamp") and value_type in ( "val_int", "val_float", "val_double", @@ -1256,6 +1257,7 @@ def get_subscriptable_field(field: Expression) -> SubscriptableReference | None: # is the right existence check. Scalar map lookups still use map_key_exists. if isinstance(base, Column) and base.column_name in TYPED_ARRAY_MAP_COLUMNS: return f.notEmpty(field) + return map_key_exists(field.parameters[0], field.parameters[1]) if isinstance(field, FunctionCall) and field.function_name in ("arrayMap", "arrayConcat"): @@ -1265,4 +1267,12 @@ def get_subscriptable_field(field: Expression) -> SubscriptableReference | None: # type_array_to_membership_array_expression_from_typed_columns). return f.notEmpty(field) + if isinstance(field, FunctionCall) and field.function_name == "cast": + # A normalized String column with an empty-string default (e.g. ai_conversation_id) + # is never NULL, so isNotNull would always be true. notEmpty distinguishes an unset + # (empty) value from a real one. + base = field.parameters[0] + if isinstance(base, Column) and base.column_name in EMPTY_STRING_DEFAULT_COLUMNS: + return f.notEmpty(field) + return f.isNotNull(field) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py index 434f2b621a..f2e852dd91 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_items/resolver_trace_item_table.py @@ -34,6 +34,7 @@ from snuba.protos.common import ( NORMALIZED_COLUMNS_EAP_ITEMS, TYPED_ARRAY_MAP_COLUMNS, + NormalizedColumn, type_array_typed_columns_select_expressions, ) from snuba.query import OrderBy, OrderByDirection, SelectedExpression @@ -189,13 +190,12 @@ def transform_expressions(expression: Expression) -> Expression: if context is None: return expression + from_column_type = NORMALIZED_COLUMNS_EAP_ITEMS.get( + context.from_column_name, + NormalizedColumn(context.from_column_name, [AttributeKey.TYPE_STRING]), + ).types[0] from_column = attribute_key_to_expression( - AttributeKey( - name=context.from_column_name, - type=NORMALIZED_COLUMNS_EAP_ITEMS.get( - context.from_column_name, [AttributeKey.TYPE_STRING] - )[0], - ) + AttributeKey(name=context.from_column_name, type=from_column_type) ) if is_existence: return get_field_existence_expression(from_column) diff --git a/tests/web/rpc/v1/test_indexed_columns.py b/tests/web/rpc/v1/test_indexed_columns.py new file mode 100644 index 0000000000..823a605b01 --- /dev/null +++ b/tests/web/rpc/v1/test_indexed_columns.py @@ -0,0 +1,200 @@ +"""EAP-620: `sentry.session_id` and `gen_ai.conversation.id` resolve to their indexed +physical columns (`session_id`, `ai_conversation_id`) as bare columns, so `=` / `IN ` +filters prune granules via the bloom-filter skip index (`bf_session_id`, +`bf_ai_conversation_id`). These tests execute a real `TraceItemTable` query, then run +`EXPLAIN indexes = 1` on the generated ClickHouse SQL to assert the skip index is applied, +and check the filter returns the right rows. +""" + +import uuid +from datetime import UTC, datetime, timedelta + +import pytest +from google.protobuf.timestamp_pb2 import Timestamp +from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import ( + Column, + TraceItemTableRequest, +) +from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta, TraceItemType +from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ( + AttributeKey, + AttributeValue, + StrArray, +) +from sentry_protos.snuba.v1.trace_item_filter_pb2 import ( + ComparisonFilter, + ExistsFilter, + TraceItemFilter, +) +from sentry_protos.snuba.v1.trace_item_pb2 import AnyValue, TraceItem + +from snuba.clusters.cluster import ClickhouseClientSettings, get_cluster +from snuba.clusters.storage_sets import StorageSetKey +from snuba.datasets.storages.factory import get_writable_storage +from snuba.datasets.storages.storage_key import StorageKey +from snuba.web.rpc.v1.endpoint_trace_item_table import EndpointTraceItemTable +from tests.helpers import write_raw_unprocessed_events + +# Two items: item A (session/conversation "a", color red), item B ("b", color blue). +SESSION_A = "11111111-1111-1111-1111-111111111111" +SESSION_B = "22222222-2222-2222-2222-222222222222" +CONV_A = "conv-a" +CONV_B = "conv-b" + + +def _item_message(ts: datetime, session_id: str, conversation_id: str, color: str) -> bytes: + item_ts = Timestamp() + item_ts.FromDatetime(ts) + received = Timestamp() + received.GetCurrentTime() + return TraceItem( + organization_id=1, + project_id=1, + item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + timestamp=item_ts, + trace_id=uuid.uuid4().hex, + item_id=uuid.uuid4().int.to_bytes(16, "little"), + received=received, + retention_days=90, + server_sample_rate=1.0, + session_id=session_id, + conversation_id=conversation_id, + attributes={ + "color": AnyValue(string_value=color), + "sentry.end_timestamp_precise": AnyValue( + double_value=(ts + timedelta(seconds=1)).timestamp() + ), + "sentry.received": AnyValue(double_value=received.seconds), + "sentry.start_timestamp_precise": AnyValue(double_value=ts.timestamp()), + "start_timestamp_ms": AnyValue(double_value=int(ts.timestamp() * 1000)), + }, + ).SerializeToString() + + +@pytest.mark.eap +@pytest.mark.redis_db +class TestIndexedColumnSkipIndex: + @pytest.fixture(autouse=True) + def setup(self, eap: None, redis_db: None) -> None: + self.base_time = datetime.now(tz=UTC).replace( + minute=0, second=0, microsecond=0 + ) - timedelta(hours=1) + self.start_ts = Timestamp(seconds=int((self.base_time - timedelta(hours=1)).timestamp())) + self.end_ts = Timestamp(seconds=int((self.base_time + timedelta(hours=2)).timestamp())) + messages = [ + _item_message(self.base_time, SESSION_A, CONV_A, "red"), + _item_message(self.base_time + timedelta(minutes=1), SESSION_B, CONV_B, "blue"), + # "green" sets neither: an unset conversation_id ingests as an empty + # ai_conversation_id, so exists on gen_ai.conversation.id must exclude it. + # (session_id is intentionally not exercised for exists: an unset session_id + # ingests as a random non-nil UUID, so absence is indistinguishable there.) + _item_message(self.base_time + timedelta(minutes=2), "", "", "green"), + ] + write_raw_unprocessed_events(get_writable_storage(StorageKey("eap_items")), messages) + + def _execute(self, filt: TraceItemFilter) -> tuple[str, list[str]]: + """Run a TraceItemTable query in debug mode; return (generated_sql, sorted colors).""" + message = TraceItemTableRequest( + meta=RequestMeta( + project_ids=[1], + organization_id=1, + cogs_category="test", + referrer="test", + start_timestamp=self.start_ts, + end_timestamp=self.end_ts, + request_id=uuid.uuid4().hex, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + debug=True, + ), + filter=filt, + columns=[Column(key=AttributeKey(type=AttributeKey.TYPE_STRING, name="color"))], + limit=100, + ) + response = EndpointTraceItemTable().execute(message) + sql = response.meta.query_info[0].metadata.sql + colors = ( + sorted(r.val_str for r in response.column_values[0].results) + if response.column_values + else [] + ) + return sql, colors + + @staticmethod + def _explain_uses_index(sql: str, index_name: str) -> bool: + conn = get_cluster(StorageSetKey.EVENTS_ANALYTICS_PLATFORM).get_query_connection( + ClickhouseClientSettings.QUERY + ) + plan = "\n".join(str(row[0]) for row in conn.execute(f"EXPLAIN indexes = 1 {sql}").results) + # The index only appears in the plan's "Skip" section when ClickHouse can analyze the + # condition against it — i.e. when the column is bare. A value-changing wrapper + # (Nullable cast, replaceAll, ...) drops it from the plan. (With a single granule + # nothing is pruned, so we assert the index is *applied to the plan*, not a count.) + return index_name in plan + + # honestly this test is probably unnecessary. It was useful when I was developing, so I figured why not merge it in. + @pytest.mark.parametrize( + "attribute_name,op,value,expected_index,expected_colors", + [ + ( + "sentry.session_id", + ComparisonFilter.OP_EQUALS, + AttributeValue(val_str=SESSION_A), + "bf_session_id", + ["red"], + ), + ( + "sentry.session_id", + ComparisonFilter.OP_IN, + AttributeValue(val_str_array=StrArray(values=[SESSION_A, SESSION_B])), + "bf_session_id", + ["blue", "red"], + ), + ( + "gen_ai.conversation.id", + ComparisonFilter.OP_EQUALS, + AttributeValue(val_str=CONV_A), + "bf_ai_conversation_id", + ["red"], + ), + ( + "gen_ai.conversation.id", + ComparisonFilter.OP_IN, + AttributeValue(val_str_array=StrArray(values=[CONV_A, CONV_B])), + "bf_ai_conversation_id", + ["blue", "red"], + ), + ], + ) + def test_filter_uses_skip_index_and_matches( + self, + attribute_name: str, + op: "ComparisonFilter.Op.ValueType", + value: AttributeValue, + expected_index: str, + expected_colors: list[str], + ) -> None: + sql, colors = self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name=attribute_name), + op=op, + value=value, + ) + ) + ) + assert colors == expected_colors + assert self._explain_uses_index(sql, expected_index), sql + + def test_exists_on_conversation_id_excludes_rows_without_a_value(self) -> None: + # gen_ai.conversation.id is optional: the "green" span never set it, so its + # ai_conversation_id is empty. exists() must treat empty as absent and exclude + # it, rather than matching every row (a non-nullable column is never NULL, so a + # plain isNotNull check would always be true). + _, colors = self._execute( + TraceItemFilter( + exists_filter=ExistsFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name="gen_ai.conversation.id") + ) + ) + ) + assert colors == ["blue", "red"]