diff --git a/snuba/clickhouse/native.py b/snuba/clickhouse/native.py index 787780f06fa..6d65d124e84 100644 --- a/snuba/clickhouse/native.py +++ b/snuba/clickhouse/native.py @@ -205,8 +205,6 @@ def query_execute() -> Any: else: result_data = query_execute() - print("result_dataaaa", result_data, with_column_types) - profile_data = ClickhouseProfile( bytes=conn.last_query.profile_info.bytes or 0, progress_bytes=conn.last_query.progress.bytes or 0, @@ -519,8 +517,6 @@ def execute( if "query_id" in settings: query_id = settings.pop("query_id") - print("isrobust???", self.__client.execute) - execute_func = ( self.__client.execute_robust if robust is True else self.__client.execute ) diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index 17498afd493..45a28af5899 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -183,11 +183,6 @@ def execute_query( robust=robust, ) - - print("readerrrr", reader) - - print("resultttt", result) - timer.mark("execute") stats.update( { diff --git a/snuba/web/rpc/common/exceptions.py b/snuba/web/rpc/common/exceptions.py index bbf23107d26..7714d25b3a1 100644 --- a/snuba/web/rpc/common/exceptions.py +++ b/snuba/web/rpc/common/exceptions.py @@ -17,10 +17,6 @@ class BadSnubaRPCRequestException(RPCRequestException): def __init__(self, message: str): super().__init__(400, message) -class OOMException(RPCRequestException): - def __init__(self, message: str): - super().__init__(241, message) - def convert_rpc_exception_to_proto( exc: Union[RPCRequestException, QueryException] diff --git a/snuba/web/rpc/v1/endpoint_time_series.py b/snuba/web/rpc/v1/endpoint_time_series.py index 5b8414b4d08..14962833734 100644 --- a/snuba/web/rpc/v1/endpoint_time_series.py +++ b/snuba/web/rpc/v1/endpoint_time_series.py @@ -2,12 +2,15 @@ import uuid from typing import Type +import sentry_sdk from sentry_protos.snuba.v1.endpoint_time_series_pb2 import ( TimeSeriesRequest, TimeSeriesResponse, ) from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType +from snuba import environment +from snuba.utils.metrics.wrapper import MetricsWrapper from snuba.web.rpc import RPCEndpoint, TraceItemDataResolver from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException from snuba.web.rpc.v1.resolvers import ResolverTimeSeries @@ -32,6 +35,8 @@ # MAX 5 minute granularity over 7 days _MAX_BUCKETS_IN_REQUEST = 2016 +metrics = MetricsWrapper(environment.metrics, "endpoint_time_series") + def _enforce_no_duplicate_labels(request: TimeSeriesRequest) -> None: labels = set() @@ -105,4 +110,10 @@ def _execute(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse: "This endpoint requires meta.trace_item_type to be set (are you requesting spans? logs?)" ) resolver = self.get_resolver(in_msg.meta.trace_item_type) - return resolver.resolve(in_msg) + try: + return resolver.resolve(in_msg) + except Exception as e: + if "DB::Exception: Memory limit (for query) exceeded" in str(e): + metrics.increment("endpoint_trace_item_table_OOM") + sentry_sdk.capture_exception(e) + raise BadSnubaRPCRequestException(str(e)) diff --git a/snuba/web/rpc/v1/endpoint_trace_item_table.py b/snuba/web/rpc/v1/endpoint_trace_item_table.py index 33a4000a137..5230d04c8d7 100644 --- a/snuba/web/rpc/v1/endpoint_trace_item_table.py +++ b/snuba/web/rpc/v1/endpoint_trace_item_table.py @@ -1,6 +1,7 @@ import uuid from typing import Type +import sentry_sdk from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import ( Column, TraceItemTableRequest, @@ -8,12 +9,16 @@ ) from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType +from snuba import environment +from snuba.utils.metrics.wrapper import MetricsWrapper from snuba.web.rpc import RPCEndpoint, TraceItemDataResolver from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException from snuba.web.rpc.v1.resolvers import ResolverTraceItemTable _GROUP_BY_DISALLOWED_COLUMNS = ["timestamp"] +metrics = MetricsWrapper(environment.metrics, "endpoint_trace_item_table") + def _apply_labels_to_columns(in_msg: TraceItemTableRequest) -> TraceItemTableRequest: def _apply_label_to_column(column: Column) -> None: @@ -104,4 +109,14 @@ def _execute(self, in_msg: TraceItemTableRequest) -> TraceItemTableResponse: "This endpoint requires meta.trace_item_type to be set (are you requesting spans? logs?)" ) resolver = self.get_resolver(in_msg.meta.trace_item_type) - return resolver.resolve(in_msg) + try: + return resolver.resolve(in_msg) + except Exception as e: + print(e.args) + print(str(e)) + print(e) + print("didthisraise??????") + if "DB::Exception: Memory limit (for query) exceeded" in str(e): + metrics.increment("endpoint_trace_item_table_OOM") + sentry_sdk.capture_exception(e) + raise BadSnubaRPCRequestException(str(e)) diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py index 79b3c44c475..428894b8fdf 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py @@ -39,7 +39,6 @@ extract_response_meta, setup_trace_query_settings, ) -from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException from snuba.web.rpc.v1.resolvers import ResolverTimeSeries from snuba.web.rpc.v1.resolvers.R_eap_spans.common.aggregation import ( ExtrapolationContext, @@ -48,14 +47,10 @@ get_confidence_interval_column, get_count_column, ) -from clickhouse_driver.errors import Error -import sentry_sdk - metrics = MetricsWrapper(environment.metrics, "endpoint_trace_item_table") - def _convert_result_timeseries( request: TimeSeriesRequest, data: list[Dict[str, Any]] ) -> Iterable[TimeSeries]: @@ -306,17 +301,11 @@ def trace_item_type(cls) -> TraceItemType.ValueType: def resolve(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse: snuba_request = _build_snuba_request(in_msg) - try: - res = run_query( - dataset=PluggableDataset(name="eap", all_entities=[]), - request=snuba_request, - timer=self._timer, - ) - except Error as e: - if e.code == 241 or "DB::Exception: Memory limit (for query) exceeded" in e.message: - metrics.increment("endpoint_trace_item_table_OOM") - sentry_sdk.capture_exception(e) - raise BadSnubaRPCRequestException(e.message) + res = run_query( + dataset=PluggableDataset(name="eap", all_entities=[]), + request=snuba_request, + timer=self._timer, + ) response_meta = extract_response_meta( in_msg.meta.request_id, diff --git a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py index 9db8ef20c38..2e9555bb5d4 100644 --- a/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py +++ b/snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py @@ -2,9 +2,7 @@ from collections import defaultdict from dataclasses import replace from typing import Any, Callable, Dict, Iterable, Sequence -import sentry_sdk -from clickhouse_driver.errors import Error from google.protobuf.json_format import MessageToDict from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import ( AggregationComparisonFilter, @@ -47,7 +45,7 @@ extract_response_meta, setup_trace_query_settings, ) -from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException, OOMException +from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException from snuba.web.rpc.v1.resolvers import ResolverTraceItemTable from snuba.web.rpc.v1.resolvers.R_eap_spans.common.aggregation import ( ExtrapolationContext, @@ -302,17 +300,11 @@ def trace_item_type(cls) -> TraceItemType.ValueType: def resolve(self, in_msg: TraceItemTableRequest) -> TraceItemTableResponse: snuba_request = _build_snuba_request(in_msg) - try: - res = run_query( - dataset=PluggableDataset(name="eap", all_entities=[]), - request=snuba_request, - timer=self._timer, - ) - except Error as e: - if e.code == 241 or "DB::Exception: Memory limit (for query) exceeded" in e.message: - metrics.increment("endpoint_trace_item_table_OOM") - sentry_sdk.capture_exception(e) - raise BadSnubaRPCRequestException(e.message) + res = run_query( + dataset=PluggableDataset(name="eap", all_entities=[]), + request=snuba_request, + timer=self._timer, + ) column_values = _convert_results(in_msg, res.result.get("data", [])) response_meta = extract_response_meta( diff --git a/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series.py b/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series.py index 016edbbaa9d..d3d336dbe1e 100644 --- a/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series.py +++ b/tests/web/rpc/v1/test_endpoint_time_series/test_endpoint_time_series.py @@ -2,8 +2,10 @@ from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any, Callable, MutableMapping +from unittest.mock import patch import pytest +from clickhouse_driver.errors import ServerException from google.protobuf.timestamp_pb2 import Timestamp from sentry_protos.snuba.v1.endpoint_time_series_pb2 import ( DataPoint, @@ -29,6 +31,7 @@ from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException +from snuba.web.rpc.v1 import endpoint_time_series from snuba.web.rpc.v1.endpoint_time_series import ( EndpointTimeSeries, _validate_time_buckets, @@ -748,6 +751,56 @@ def test_with_non_existent_attribute(self) -> None: ) ] + def test_OOM(self) -> None: + ts = Timestamp() + ts.GetCurrentTime() + tstart = Timestamp(seconds=ts.seconds - 3600) + message = TimeSeriesRequest( + meta=RequestMeta( + project_ids=[1, 2, 3], + organization_id=1, + cogs_category="something", + referrer="something", + start_timestamp=tstart, + end_timestamp=ts, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + ), + aggregations=[ + AttributeAggregation( + aggregate=Function.FUNCTION_AVG, + key=AttributeKey( + type=AttributeKey.TYPE_FLOAT, name="sentry.duration" + ), + label="p50", + ), + AttributeAggregation( + aggregate=Function.FUNCTION_P95, + key=AttributeKey( + type=AttributeKey.TYPE_FLOAT, name="sentry.duration" + ), + label="p90", + ), + ], + granularity_secs=60, + ) + + with patch( + "clickhouse_driver.client.Client.execute", + side_effect=ServerException( + "DB::Exception: Received from snuba-events-analytics-platform-1-1:1111. DB::Exception: Memory limit (for query) exceeded: would use 1.11GiB (attempt to allocate chunk of 111111 bytes), maximum: 1.11 GiB. Blahblahblahblahblahblahblah", + code=241, + ), + ), patch.object( + endpoint_time_series.metrics, "increment" + ) as metrics_mock, patch( + "snuba.web.rpc.v1.endpoint_trace_item_table.sentry_sdk.capture_exception" + ) as sentry_sdk_mock: + with pytest.raises(BadSnubaRPCRequestException) as e: + EndpointTimeSeries().execute(message) + assert "DB::Exception: Memory limit (for query) exceeded" in str(e.value) + metrics_mock.assert_called_once_with("endpoint_trace_item_table_OOM") + sentry_sdk_mock.assert_called_once() + class TestUtils: def test_no_duplicate_labels(self) -> None: diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py index 5775d458770..4bae940d0ea 100644 --- a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py +++ b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py @@ -2,8 +2,10 @@ import uuid from datetime import datetime, timedelta from typing import Any, Mapping +from unittest.mock import patch import pytest +from clickhouse_driver.errors import ServerException from google.protobuf.json_format import MessageToDict, ParseDict from google.protobuf.timestamp_pb2 import Timestamp from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import ( @@ -41,6 +43,7 @@ from snuba.datasets.storages.factory import get_storage from snuba.datasets.storages.storage_key import StorageKey from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException +from snuba.web.rpc.v1 import endpoint_trace_item_table from snuba.web.rpc.v1.endpoint_trace_item_table import ( EndpointTraceItemTable, _apply_labels_to_columns, @@ -185,6 +188,53 @@ def test_basic(self) -> None: error_proto.ParseFromString(response.data) assert response.status_code == 200, error_proto + def test_OOM(self) -> None: + ts = Timestamp() + ts.GetCurrentTime() + message = TraceItemTableRequest( + meta=RequestMeta( + project_ids=[1, 2, 3], + organization_id=1, + cogs_category="something", + referrer="something", + start_timestamp=ts, + end_timestamp=ts, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + ), + filter=TraceItemFilter( + exists_filter=ExistsFilter( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name="color") + ) + ), + columns=[ + Column(key=AttributeKey(type=AttributeKey.TYPE_STRING, name="location")) + ], + order_by=[ + TraceItemTableRequest.OrderBy( + column=Column( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name="location") + ) + ) + ], + limit=10, + ) + with patch( + "clickhouse_driver.client.Client.execute", + side_effect=ServerException( + "DB::Exception: Received from snuba-events-analytics-platform-1-1:1111. DB::Exception: Memory limit (for query) exceeded: would use 1.11GiB (attempt to allocate chunk of 111111 bytes), maximum: 1.11 GiB. Blahblahblahblahblahblahblah", + code=241, + ), + ), patch.object( + endpoint_trace_item_table.metrics, "increment" + ) as metrics_mock, patch( + "snuba.web.rpc.v1.endpoint_trace_item_table.sentry_sdk.capture_exception" + ) as sentry_sdk_mock: + with pytest.raises(BadSnubaRPCRequestException) as e: + EndpointTraceItemTable().execute(message) + assert "DB::Exception: Memory limit (for query) exceeded" in str(e.value) + metrics_mock.assert_called_once_with("endpoint_trace_item_table_OOM") + sentry_sdk_mock.assert_called_once() + def test_errors_without_type(self) -> None: ts = Timestamp() ts.GetCurrentTime() @@ -269,7 +319,6 @@ def test_with_data(self, setup_teardown: Any) -> None: meta=ResponseMeta(request_id="be3123b3-2e5d-4eb9-bb48-f38eaa9e8480"), ) assert response == expected_response - assert False def test_booleans_and_number_compares_backward_compat( self, setup_teardown: Any