rpc base class

getsentry · Feb 3, 2025 · 2c8d44e · 2c8d44e
1 parent 386cb84
commit 2c8d44e
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 23 deletions.
diff --git a/snuba/web/rpc/__init__.py b/snuba/web/rpc/__init__.py
@@ -119,6 +119,15 @@ def execute(self, in_msg: Tin) -> Tout:
         error = None
         try:
             out = self._execute(in_msg)
+        except QueryException as e:
+            if (
+                "error_code" in e.extra["stats"]
+                and e.extra["stats"]["error_code"] == 241
+            ):
+                self.metrics.increment("OOM_query")
+                sentry_sdk.capture_exception(e)
+            out = self.response_class()()
+            error = e
         except Exception as e:
             out = self.response_class()()
             error = e
@@ -198,12 +207,7 @@ def run_rpc_handler(
 
     try:
         return cast(ProtobufMessage, endpoint.execute(deserialized_protobuf))
-    except RPCRequestException as e:
-        return convert_rpc_exception_to_proto(e)
-    except QueryException as e:
-        if "error_code" in e.extra["stats"] and e.extra["stats"]["error_code"] == 241:
-            endpoint.metrics.increment("OOM_query")
-            sentry_sdk.capture_exception(e)
+    except (RPCRequestException, QueryException) as e:
         return convert_rpc_exception_to_proto(e)
     except Exception as e:
         sentry_sdk.capture_exception(e)

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
@@ -13,7 +13,6 @@
     TimeSeriesRequest,
 )
 from sentry_protos.snuba.v1.error_pb2 import Error
-from sentry_protos.snuba.v1.error_pb2 import Error as ErrorProto
 from sentry_protos.snuba.v1.request_common_pb2 import RequestMeta, TraceItemType
 from sentry_protos.snuba.v1.trace_item_attribute_pb2 import (
     AttributeAggregation,
@@ -31,7 +30,8 @@
 
 from snuba.datasets.storages.factory import get_storage
 from snuba.datasets.storages.storage_key import StorageKey
-from snuba.web.rpc import RPCEndpoint, run_rpc_handler
+from snuba.web import QueryException
+from snuba.web.rpc import RPCEndpoint
 from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
 from snuba.web.rpc.v1.endpoint_time_series import (
     EndpointTimeSeries,
@@ -794,13 +794,9 @@ def test_OOM(self, monkeypatch: Any) -> None:
                 code=241,
             ),
         ), patch("snuba.web.rpc.sentry_sdk.capture_exception") as sentry_sdk_mock:
-            resp = run_rpc_handler(
-                "EndpointTimeSeries", "v1", message.SerializeToString()
-            )
-            assert isinstance(resp, ErrorProto)
-            assert "DB::Exception: Memory limit (for query) exceeded" in str(
-                resp.message
-            )
+            with pytest.raises(QueryException) as e:
+                EndpointTimeSeries().execute(message)
+            assert "DB::Exception: Memory limit (for query) exceeded" in str(e.value)
 
             sentry_sdk_mock.assert_called_once()
             assert metrics_mock.increment.call_args_list.count(call("OOM_query")) == 1

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
@@ -42,7 +42,8 @@
 
 from snuba.datasets.storages.factory import get_storage
 from snuba.datasets.storages.storage_key import StorageKey
-from snuba.web.rpc import RPCEndpoint, run_rpc_handler
+from snuba.web import QueryException
+from snuba.web.rpc import RPCEndpoint
 from snuba.web.rpc.common.exceptions import BadSnubaRPCRequestException
 from snuba.web.rpc.v1.endpoint_trace_item_table import (
     EndpointTraceItemTable,
@@ -227,13 +228,9 @@ def test_OOM(self, monkeypatch: Any) -> None:
                 code=241,
             ),
         ), patch("snuba.web.rpc.sentry_sdk.capture_exception") as sentry_sdk_mock:
-            resp = run_rpc_handler(
-                "EndpointTraceItemTable", "v1", message.SerializeToString()
-            )
-            assert isinstance(resp, ErrorProto)
-            assert "DB::Exception: Memory limit (for query) exceeded" in str(
-                resp.message
-            )
+            with pytest.raises(QueryException) as e:
+                EndpointTraceItemTable().execute(message)
+            assert "DB::Exception: Memory limit (for query) exceeded" in str(e.value)
 
             sentry_sdk_mock.assert_called_once()
             assert metrics_mock.increment.call_args_list.count(call("OOM_query")) == 1