diff --git a/snuba/clickhouse/error_codes.py b/snuba/clickhouse/error_codes.py index 59668bc712..eee37a9866 100644 --- a/snuba/clickhouse/error_codes.py +++ b/snuba/clickhouse/error_codes.py @@ -11,6 +11,7 @@ class ErrorCodes(IntEnum): + CANNOT_COMPILE_REGEXP = 36 ILLEGAL_TYPE_OF_ARGUMENT = 43 ILLEGAL_COLUMN = 44 UNKNOWN_FUNCTION = 46 diff --git a/snuba/querylog/query_metadata.py b/snuba/querylog/query_metadata.py index cb91ce4313..5e97475e44 100644 --- a/snuba/querylog/query_metadata.py +++ b/snuba/querylog/query_metadata.py @@ -123,6 +123,7 @@ def slo(self) -> SLO: ErrorCodes.ILLEGAL_AGGREGATION: RequestStatus.INVALID_REQUEST, ErrorCodes.TOO_MANY_SIMULTANEOUS_QUERIES: RequestStatus.CLICKHOUSE_MAX_QUERIES_EXCEEDED, ErrorCodes.CANNOT_PARSE_DOMAIN_VALUE_FROM_STRING: RequestStatus.INVALID_REQUEST, + ErrorCodes.CANNOT_COMPILE_REGEXP: RequestStatus.INVALID_REQUEST, } diff --git a/snuba/web/rpc/common/common.py b/snuba/web/rpc/common/common.py index f5da7b58f2..f064766f07 100644 --- a/snuba/web/rpc/common/common.py +++ b/snuba/web/rpc/common/common.py @@ -567,7 +567,9 @@ def _attribute_value_to_expression(v: AttributeValue) -> Expression: AnyAttributeFilter.OP_NOT_IN, } -_POSITIVE_OP_FOR_NEGATIVE: dict[int, int] = { +_POSITIVE_OP_FOR_NEGATIVE: dict[ + AnyAttributeFilter.Op.ValueType, AnyAttributeFilter.Op.ValueType +] = { AnyAttributeFilter.OP_NOT_EQUALS: AnyAttributeFilter.OP_EQUALS, AnyAttributeFilter.OP_NOT_LIKE: AnyAttributeFilter.OP_LIKE, AnyAttributeFilter.OP_NOT_IN: AnyAttributeFilter.OP_IN, @@ -607,15 +609,15 @@ def _array_value_length(v: AttributeValue, value_type: str) -> int: def _validate_comparison_filter_type_array( op: ComparisonFilter.Op.ValueType, v: AttributeValue, key: AttributeKey ) -> None: - if op in (ComparisonFilter.OP_LIKE, ComparisonFilter.OP_NOT_LIKE): + if op in (ComparisonFilter.OP_LIKE, ComparisonFilter.OP_NOT_LIKE, ComparisonFilter.OP_REGEXP): if v.WhichOneof("value") != "val_str": - raise BadSnubaRPCRequestException( - "LIKE/NOT_LIKE on array keys requires a string pattern" - ) - # LIKE only matches string elements, so it makes sense only for string arrays. + label = "REGEXP" if op == ComparisonFilter.OP_REGEXP else "LIKE/NOT_LIKE" + raise BadSnubaRPCRequestException(f"{label} on array keys requires a string pattern") + # LIKE/REGEXP only match string elements, so they make sense only for string arrays. if key.type not in (AttributeKey.Type.TYPE_ARRAY, AttributeKey.Type.TYPE_ARRAY_STRING): + label = "REGEXP" if op == ComparisonFilter.OP_REGEXP else "LIKE/NOT_LIKE" raise BadSnubaRPCRequestException( - "LIKE/NOT_LIKE on array keys is only supported on string arrays " + f"{label} on array keys is only supported on string arrays " f"(TYPE_ARRAY_STRING), got {AttributeKey.Type.Name(key.type)}" ) return @@ -854,6 +856,57 @@ def _typed_array_like_expression( ) +def _regexp_match(value: Expression, pattern: Expression, ignore_case: bool) -> FunctionCall: + # Needs ClickHouse > 26.8 for matchCaseInsensitive; lower() is a stand-in until then. + if ignore_case: + return f.match(f.lower(value), f.lower(pattern)) + return f.match(value, pattern) + + +def _is_valid_regexp_pattern(v: AttributeValue) -> None: + if v.WhichOneof("value") != "val_str" or v.val_str == "": + raise BadSnubaRPCRequestException("REGEXP pattern must be a non-empty string") + + +def _any_attribute_op_expression( + *, + filter_: AnyAttributeFilter, + op: AnyAttributeFilter.Op.ValueType, + element: Argument, + value_expr: Expression, + membership_as_has: bool, +) -> Expression: + match op: + case AnyAttributeFilter.OP_EQUALS: + if filter_.ignore_case: + return f.equals(f.lower(element), f.lower(value_expr)) + return f.equals(element, value_expr) + case AnyAttributeFilter.OP_LIKE: + if filter_.ignore_case: + return f.ilike(element, value_expr) + return f.like(element, value_expr) + case AnyAttributeFilter.OP_REGEXP: + return _regexp_match(element, value_expr, filter_.ignore_case) + case AnyAttributeFilter.OP_IN: + if filter_.ignore_case: + if filter_.value.WhichOneof("value") == "val_str_array": + lowered = [literal(s.lower()) for s in filter_.value.val_str_array.values] + else: + lowered = [ + literal(elem.val_str.lower()) for elem in filter_.value.val_array.values + ] + return _in_or_has( + f.lower(element), + literals_array(None, lowered), + as_has=membership_as_has, + ) + return _in_or_has(element, value_expr, as_has=membership_as_has) + case _: + raise BadSnubaRPCRequestException( + f"Unsupported any_attribute_filter op: {AnyAttributeFilter.Op.Name(filter_.op)}" + ) + + def _any_attribute_filter_to_expression( filt: AnyAttributeFilter, *, @@ -918,11 +971,16 @@ def _any_attribute_filter_to_expression( ) col_name = _VALUE_TYPE_TO_COLUMN[value_type] - # LIKE/NOT_LIKE only makes sense on string columns - if effective_op == AnyAttributeFilter.OP_LIKE and col_name not in _STRING_COLUMNS: - raise BadSnubaRPCRequestException( - "LIKE/NOT_LIKE operations are only supported on string values" - ) + # LIKE/NOT_LIKE/REGEXP only makes sense on string columns + if ( + effective_op in (AnyAttributeFilter.OP_LIKE, AnyAttributeFilter.OP_REGEXP) + and col_name not in _STRING_COLUMNS + ): + label = "REGEXP" if effective_op == AnyAttributeFilter.OP_REGEXP else "LIKE/NOT_LIKE" + raise BadSnubaRPCRequestException(f"{label} operations are only supported on string values") + + if effective_op == AnyAttributeFilter.OP_REGEXP: + _is_valid_regexp_pattern(v) # ignore_case uses lower() which only works on string columns if filt.ignore_case and col_name not in _STRING_COLUMNS: @@ -932,35 +990,13 @@ def _any_attribute_filter_to_expression( # 3. Build the lambda comparison x = Argument(None, "x") - - if effective_op == AnyAttributeFilter.OP_EQUALS: - if filt.ignore_case: - comparison = f.equals(f.lower(x), f.lower(v_expression)) - else: - comparison = f.equals(x, v_expression) - elif effective_op == AnyAttributeFilter.OP_LIKE: - if filt.ignore_case: - comparison = f.ilike(x, v_expression) - else: - comparison = f.like(x, v_expression) - elif effective_op == AnyAttributeFilter.OP_IN: - if filt.ignore_case: - if value_type == "val_str_array": - lowered = [literal(s.lower()) for s in v.val_str_array.values] - else: - lowered = [literal(elem.val_str.lower()) for elem in v.val_array.values] - comparison = _in_or_has( - f.lower(x), - literals_array(None, lowered), - as_has=membership_as_has, - ) - else: - comparison = _in_or_has(x, v_expression, as_has=membership_as_has) - else: - raise BadSnubaRPCRequestException( - f"Unsupported any_attribute_filter op: {AnyAttributeFilter.Op.Name(filt.op)}" - ) - + comparison = _any_attribute_op_expression( + filter_=filt, + op=effective_op, + element=x, + value_expr=v_expression, + membership_as_has=membership_as_has, + ) lam = Lambda(None, ("x",), comparison) # 4. Build the arrayExists expression for the single matching column. @@ -1261,6 +1297,26 @@ def trace_item_filters_to_expression( value, exists = _map_backed_operands(k) return and_cond(exists, comparison_function(value, v_expression)) return comparison_function(k_expression, v_expression) + if op == ComparisonFilter.OP_REGEXP: + _is_valid_regexp_pattern(v) + ignore_case = item_filter.comparison_filter.ignore_case + if k.type in ARRAY_TYPES: + return f.arrayExists( + Lambda( + None, + ("x",), + _regexp_match(Argument(None, "x"), v_expression, ignore_case), + ), + type_array_typed_column_native_array(k, "attributes_array_string"), + ) + if k.type != AttributeKey.Type.TYPE_STRING: + raise BadSnubaRPCRequestException( + "the REGEXP comparison is only supported on string and array keys" + ) + if _is_map_backed_key(k): + value, exists = _map_backed_operands(k) + return and_cond(exists, _regexp_match(value, v_expression, ignore_case)) + return _regexp_match(k_expression, v_expression, ignore_case) if op == ComparisonFilter.OP_NOT_LIKE: if k.type in ARRAY_TYPES: return not_cond( diff --git a/snuba/web/rpc/common/exceptions.py b/snuba/web/rpc/common/exceptions.py index fe6c503468..bc3b58dda1 100644 --- a/snuba/web/rpc/common/exceptions.py +++ b/snuba/web/rpc/common/exceptions.py @@ -3,6 +3,7 @@ from google.protobuf import any_pb2, struct_pb2 from sentry_protos.snuba.v1.error_pb2 import Error as ErrorProto +from snuba.clickhouse.error_codes import ErrorCodes from snuba.web import QueryException @@ -61,5 +62,9 @@ def convert_rpc_exception_to_proto(exc: RPCRequestException | QueryException) -> inferred_status = 500 if exc.exception_type == "RateLimitExceeded": inferred_status = 429 + else: + error_code = exc.extra.get("stats", {}).get("error_code") + if error_code == ErrorCodes.CANNOT_COMPILE_REGEXP: + inferred_status = 400 return ErrorProto(code=inferred_status, message=str(exc)) diff --git a/tests/querylog/test_query_metadata.py b/tests/querylog/test_query_metadata.py index d42acb0611..c4f837e4a5 100644 --- a/tests/querylog/test_query_metadata.py +++ b/tests/querylog/test_query_metadata.py @@ -79,3 +79,9 @@ def test_get_request_status_timeout(self) -> None: status = get_request_status(error) assert status.status == RequestStatus.CLICKHOUSE_TIMEOUT assert status.slo == SLO.AGAINST + + def test_get_request_status_cannot_compile_regexp(self) -> None: + error = ClickhouseError("cannot compile regexp", code=ErrorCodes.CANNOT_COMPILE_REGEXP) + status = get_request_status(error) + assert status.status == RequestStatus.INVALID_REQUEST + assert status.slo == SLO.FOR diff --git a/tests/web/rpc/test_common.py b/tests/web/rpc/test_common.py index d8d11c6b78..f27948b806 100644 --- a/tests/web/rpc/test_common.py +++ b/tests/web/rpc/test_common.py @@ -29,9 +29,10 @@ ExistsFilter, TraceItemFilter, ) -from sentry_protos.snuba.v1.trace_item_pb2 import AnyValue +from sentry_protos.snuba.v1.trace_item_pb2 import AnyValue, ArrayValue from snuba import settings +from snuba.clickhouse.error_codes import ErrorCodes from snuba.datasets.storages.factory import get_writable_storage from snuba.datasets.storages.storage_key import StorageKey from snuba.protos.common import ( @@ -53,6 +54,7 @@ SubscriptableReference, ) from snuba.query.logical import Query +from snuba.web import QueryException from snuba.web.rpc.common.common import ( USE_INDEXED_NAME_ORGANIZATION_IDS_OPTION, _any_attribute_filter_to_expression, @@ -1198,6 +1200,20 @@ def test_convert_rpc_exception_to_proto_packs_details() -> None: assert unpacked == routing_decision_log_dict +def test_convert_rpc_exception_cannot_compile_regexp_is_400() -> None: + exc = QueryException.from_args( + "ClickhouseError", + "Code: 36. Cannot compile regexp", + extra={ + "stats": {"error_code": ErrorCodes.CANNOT_COMPILE_REGEXP}, + "sql": "SELECT 1", + "experiments": {}, + }, + ) + proto = convert_rpc_exception_to_proto(exc) + assert proto.code == 400 + + class TestAnyAttributeFilter: def test_like_on_non_string_value_raises(self) -> None: filt = AnyAttributeFilter( @@ -1447,6 +1463,53 @@ def test_like_finds_target_span(self) -> None: ) assert colors == ["red"] + def test_regexp_finds_target_span(self) -> None: + colors = self._execute( + TraceItemFilter( + any_attribute_filter=AnyAttributeFilter( + op=AnyAttributeFilter.OP_REGEXP, + value=AttributeValue(val_str=self.UNIQUE_VALUE), + ) + ) + ) + assert colors == ["red"] + + def test_regexp_ignore_case(self) -> None: + colors = self._execute( + TraceItemFilter( + any_attribute_filter=AnyAttributeFilter( + op=AnyAttributeFilter.OP_REGEXP, + value=AttributeValue(val_str=self.UNIQUE_VALUE.upper()), + ignore_case=True, + ) + ) + ) + assert colors == ["red"] + + def test_regexp_empty_pattern_raises(self) -> None: + with pytest.raises( + BadSnubaRPCRequestException, match="REGEXP pattern must be a non-empty string" + ): + self._execute( + TraceItemFilter( + any_attribute_filter=AnyAttributeFilter( + op=AnyAttributeFilter.OP_REGEXP, + value=AttributeValue(val_str=""), + ) + ) + ) + + def test_regexp_on_non_string_value_raises(self) -> None: + with pytest.raises(BadSnubaRPCRequestException, match="REGEXP"): + self._execute( + TraceItemFilter( + any_attribute_filter=AnyAttributeFilter( + op=AnyAttributeFilter.OP_REGEXP, + value=AttributeValue(val_int=42), + ) + ) + ) + def test_not_equals_excludes_target_span(self) -> None: """OP_NOT_EQUALS on the unique value should return the other two spans.""" colors = self._execute( @@ -1640,11 +1703,201 @@ def test_like_wildcard_matches_present_not_absent(self) -> None: # absent key. assert self._execute(ComparisonFilter.OP_LIKE, value="%") == ["blue", "red"] + def test_regexp_dotstar_matches_present_not_absent(self) -> None: + assert self._execute(ComparisonFilter.OP_REGEXP, value=".*") == ["blue", "red"] + def test_not_like_wildcard_matches_only_absent(self) -> None: # Present rows all `like '%'`, so only the absent key survives NOT LIKE. assert self._execute(ComparisonFilter.OP_NOT_LIKE, value="%") == ["green"] +@pytest.mark.eap +@pytest.mark.redis_db +class TestRegexpComparisonFilterIntegration: + """Ingest spans, then REGEXP-filter them through EndpointTraceItemTable.""" + + ATTR = "test.regexp.haystack" + ARRAY_ATTR = "test.regexp.tags" + BATCH_ATTR = "test.regexp.batch" + + @pytest.fixture(autouse=True) + def setup(self, eap: None, redis_db: None) -> None: + self.batch = f"batch-{uuid.uuid4().hex}" + 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())) + batch = AnyValue(string_value=self.batch) + messages = [ + gen_item_message( + start_timestamp=self.base_time, + attributes={ + self.BATCH_ATTR: batch, + self.ATTR: AnyValue(string_value="db.query"), + self.ARRAY_ATTR: AnyValue( + array_value=ArrayValue( + values=[ + AnyValue(string_value="timeout"), + AnyValue(string_value="retry"), + ] + ) + ), + "color": AnyValue(string_value="red"), + }, + ), + gen_item_message( + start_timestamp=self.base_time + timedelta(minutes=1), + attributes={ + self.BATCH_ATTR: batch, + self.ATTR: AnyValue(string_value="http.server"), + self.ARRAY_ATTR: AnyValue( + array_value=ArrayValue(values=[AnyValue(string_value="ok")]) + ), + "color": AnyValue(string_value="blue"), + }, + ), + ] + storage = get_writable_storage(StorageKey("eap_items")) + write_raw_unprocessed_events(storage, messages) + + def _execute( + self, + filt: TraceItemFilter, + ) -> list[str]: + 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, + ), + filter=TraceItemFilter( + and_filter=AndFilter( + filters=[ + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey( + type=AttributeKey.TYPE_STRING, name=self.BATCH_ATTR + ), + op=ComparisonFilter.OP_EQUALS, + value=AttributeValue(val_str=self.batch), + ) + ), + filt, + ] + ) + ), + columns=[Column(key=AttributeKey(type=AttributeKey.TYPE_STRING, name="color"))], + order_by=[ + TraceItemTableRequest.OrderBy( + column=Column(key=AttributeKey(type=AttributeKey.TYPE_STRING, name="color")) + ) + ], + limit=100, + ) + response = EndpointTraceItemTable().execute(message) + if not response.column_values: + return [] + return sorted(r.val_str for r in response.column_values[0].results) + + def test_regexp_on_string_key(self) -> None: + assert self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name=self.ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_str=r"db\..*"), + ) + ) + ) == ["red"] + + def test_regexp_on_string_key_ignore_case(self) -> None: + assert self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name=self.ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_str=r"DB\..*"), + ignore_case=True, + ) + ) + ) == ["red"] + + def test_regexp_on_array_key(self) -> None: + assert self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_ARRAY, name=self.ARRAY_ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_str="time.*"), + ) + ) + ) == ["red"] + + def test_regexp_on_array_key_ignore_case(self) -> None: + assert self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_ARRAY, name=self.ARRAY_ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_str="TIME.*"), + ignore_case=True, + ) + ) + ) == ["red"] + + def test_regexp_empty_pattern_raises(self) -> None: + with pytest.raises( + BadSnubaRPCRequestException, match="REGEXP pattern must be a non-empty string" + ): + self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name=self.ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_str=""), + ) + ) + ) + + def test_regexp_non_string_pattern_raises(self) -> None: + with pytest.raises( + BadSnubaRPCRequestException, match="REGEXP pattern must be a non-empty string" + ): + self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name=self.ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_double=1.5), + ) + ) + ) + + def test_regexp_pattern_is_not_interpolated_into_sql(self) -> None: + # If PATTERN were concatenated into match(col, 'PATTERN'), `' --` closes + # the SQL string and comments out the rest of the statement (including + # match's closing paren) — a syntax-error / DoS injection. A dump-all + # payload like `') OR 1 --` needs a raw `)`, which RE2 rejects as an + # unmatched paren, so that vector isn't available on a bound pattern. + # Bound, this is just the regexp `db.query|' --`, which still matches + # the seeded row (proves ingest ran). Interpolated, the query fails. + assert self._execute( + TraceItemFilter( + comparison_filter=ComparisonFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name=self.ATTR), + op=ComparisonFilter.OP_REGEXP, + value=AttributeValue(val_str="db.query|' --"), + ) + ) + ) == ["red"] + + class TestSemverSortKey: def test_expression_structure(self) -> None: expr = semver_sort_key(column("release"))