error code

snuba/web/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Mapping, TypedDict, cast
+from typing import Any, Mapping, NotRequired, TypedDict, cast
 
 from snuba.reader import Column, Result, Row, transform_rows
 from snuba.utils.serializable_exception import JsonSerializable, SerializableException
@@ -11,6 +11,7 @@ class QueryExtraData(TypedDict):
     stats: Mapping[str, Any]
     sql: str
     experiments: Mapping[str, Any]
+    code: NotRequired[int]
 
 
 class QueryException(SerializableException):

snuba/web/db_query.py

@@ -509,6 +509,17 @@ def _raw_query(
             error_code=error_code,
             triggered_rate_limiter=str(trigger_rate_limiter),
         )
+        if error_code == 241:
+            raise QueryException.from_args(
+                cause.__class__.__name__,
+                str(cause),
+                {
+                    "stats": stats,
+                    "sql": sql,
+                    "experiments": clickhouse_query.get_experiments(),
+                    "code": 241,
+                },
+            ) from cause
         raise QueryException.from_args(
             # This exception needs to have the message of the cause in it for sentry
             # to pick it up properly

snuba/web/rpc/v1/endpoint_time_series.py

@@ -11,6 +11,7 @@
 
 from snuba import environment
 from snuba.utils.metrics.wrapper import MetricsWrapper
+from snuba.web import QueryException
 from snuba.web.rpc import RPCEndpoint, TraceItemDataResolver
 from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
 from snuba.web.rpc.v1.resolvers import ResolverTimeSeries
@@ -112,8 +113,10 @@ def _execute(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse:
         resolver = self.get_resolver(in_msg.meta.trace_item_type)
         try:
             return resolver.resolve(in_msg)
-        except Exception as e:
-            if "DB::Exception: Memory limit (for query) exceeded" in str(e):
+        except QueryException as e:
+            if e.extra[
+                "code"
+            ] == 241 or "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))

snuba/web/rpc/v1/endpoint_trace_item_table.py

@@ -11,6 +11,7 @@
 
 from snuba import environment
 from snuba.utils.metrics.wrapper import MetricsWrapper
+from snuba.web import QueryException
 from snuba.web.rpc import RPCEndpoint, TraceItemDataResolver
 from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
 from snuba.web.rpc.v1.resolvers import ResolverTraceItemTable
@@ -111,8 +112,10 @@ def _execute(self, in_msg: TraceItemTableRequest) -> TraceItemTableResponse:
         resolver = self.get_resolver(in_msg.meta.trace_item_type)
         try:
             return resolver.resolve(in_msg)
-        except Exception as e:
-            if "DB::Exception: Memory limit (for query) exceeded" in str(e):
+        except QueryException as e:
+            if e.extra[
+                "code"
+            ] == 241 or "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))

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

@@ -302,7 +302,6 @@ def resolve(self, in_msg: TimeSeriesRequest) -> TimeSeriesResponse:
             request=snuba_request,
             timer=self._timer,
         )
-
         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

@@ -301,7 +301,6 @@ def resolve(self, in_msg: TraceItemTableRequest) -> TraceItemTableResponse:
             request=snuba_request,
             timer=self._timer,
         )
-
         column_values = _convert_results(in_msg, res.result.get("data", []))
         response_meta = extract_response_meta(
             in_msg.meta.request_id,