track OOM

snuba/clickhouse/native.py

@@ -205,6 +205,8 @@ 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,
@@ -517,6 +519,8 @@ 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
         )

snuba/web/db_query.py

@@ -183,6 +183,10 @@ def execute_query(
         robust=robust,
     )
 
+    print("readerrrr", reader)
+
+    print("resultttt", result)
+
     timer.mark("execute")
     stats.update(
         {

snuba/web/rpc/common/exceptions.py

@@ -18,6 +18,11 @@ 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]
 ) -> ErrorProto:

snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_time_series.py

@@ -3,6 +3,8 @@
 from datetime import datetime
 from typing import Any, Dict, Iterable
 
+import sentry_sdk
+from clickhouse_driver.errors import Error
 from google.protobuf.json_format import MessageToDict
 from google.protobuf.timestamp_pb2 import Timestamp
 from sentry_protos.snuba.v1.endpoint_time_series_pb2 import (
@@ -14,6 +16,7 @@
 from sentry_protos.snuba.v1.request_common_pb2 import TraceItemType
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import ExtrapolationMode
 
+from snuba import environment
 from snuba.attribution.appid import AppID
 from snuba.attribution.attribution_info import AttributionInfo
 from snuba.datasets.entities.entity_key import EntityKey
@@ -26,6 +29,7 @@
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
+from snuba.utils.metrics.wrapper import MetricsWrapper
 from snuba.web.query import run_query
 from snuba.web.rpc.common.common import (
     attribute_key_to_expression,
@@ -37,6 +41,7 @@
     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,
@@ -46,6 +51,8 @@
     get_count_column,
 )
 
+metrics = MetricsWrapper(environment.metrics, "endpoint_trace_item_table")
+
 
 def _convert_result_timeseries(
     request: TimeSeriesRequest, data: list[Dict[str, Any]]
@@ -297,11 +304,21 @@ def trace_item_type(cls) -> TraceItemType.ValueType:
 
     def resolve(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse:
         snuba_request = _build_snuba_request(in_msg)
-        res = run_query(
-            dataset=PluggableDataset(name="eap", all_entities=[]),
-            request=snuba_request,
-            timer=self._timer,
-        )
+        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)
+
         response_meta = extract_response_meta(
             in_msg.meta.request_id,
             in_msg.meta.debug,

snuba/web/rpc/v1/resolvers/R_eap_spans/resolver_trace_item_table.py

@@ -3,6 +3,8 @@
 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,
@@ -18,6 +20,7 @@
     ExtrapolationMode,
 )
 
+from snuba import environment
 from snuba.attribution.appid import AppID
 from snuba.attribution.attribution_info import AttributionInfo
 from snuba.datasets.entities.entity_key import EntityKey
@@ -31,6 +34,7 @@
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request as SnubaRequest
+from snuba.utils.metrics.wrapper import MetricsWrapper
 from snuba.web.query import run_query
 from snuba.web.rpc.common.common import (
     apply_virtual_columns,
@@ -43,7 +47,7 @@
     extract_response_meta,
     setup_trace_query_settings,
 )
-from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
+from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException, OOMException
 from snuba.web.rpc.v1.resolvers import ResolverTraceItemTable
 from snuba.web.rpc.v1.resolvers.R_eap_spans.common.aggregation import (
     ExtrapolationContext,
@@ -55,6 +59,8 @@
 
 _DEFAULT_ROW_LIMIT = 10_000
 
+metrics = MetricsWrapper(environment.metrics, "endpoint_trace_item_table")
+
 
 def aggregation_filter_to_expression(agg_filter: AggregationFilter) -> Expression:
     op_to_expr = {
@@ -296,11 +302,21 @@ def trace_item_type(cls) -> TraceItemType.ValueType:
 
     def resolve(self, in_msg: TraceItemTableRequest) -> TraceItemTableResponse:
         snuba_request = _build_snuba_request(in_msg)
-        res = run_query(
-            dataset=PluggableDataset(name="eap", all_entities=[]),
-            request=snuba_request,
-            timer=self._timer,
-        )
+        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)
+
         column_values = _convert_results(in_msg, res.result.get("data", []))
         response_meta = extract_response_meta(
             in_msg.meta.request_id,

tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table.py

@@ -269,6 +269,7 @@ 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