error code

getsentry · Jan 31, 2025 · 053bda5 · 053bda5
1 parent e55a1a2
commit 053bda5
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py
@@ -484,6 +484,17 @@ def _raw_query(
             trigger_rate_limiter = cause.extra_data.get("scope", "")
         elif isinstance(cause, ClickhouseError):
             error_code = cause.code
+            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
             status = get_query_status_from_error_codes(error_code)
 
             with configure_scope() as scope:

diff --git a/snuba/web/rpc/v1/endpoint_time_series.py b/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_data["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))
diff --git a/snuba/web/rpc/v1/endpoint_trace_item_table.py b/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_data["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))