fix: Catch and surface BaseException()

mukund-ananthu · mukund-ananthu · commit d43c2e3188e6 · 2024-03-05T15:55:01.000Z
Currently, the behavior of the library is to catch Exception() when encountered in the user provided callback, surface it to the calling code and shut down the client.

However, for BaseException() encountered in the user provided callback, the BaseException() is not surfaced to the calling code and the client is not shut down.

Make the behavior of the client when BaseException() is encountered consistent with the behavior of the client when Exception() is encountered.
diff --git a/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
@@ -112,7 +112,7 @@ def _wrap_as_exception(maybe_exception: Any) -> BaseException:
 
 def _wrap_callback_errors(
     callback: Callable[["google.cloud.pubsub_v1.subscriber.message.Message"], Any],
-    on_callback_error: Callable[[Exception], Any],
+    on_callback_error: Callable[[BaseException], Any],
     message: "google.cloud.pubsub_v1.subscriber.message.Message",
 ):
     """Wraps a user callback so that if an exception occurs the message is
@@ -124,7 +124,7 @@ def _wrap_callback_errors(
     """
     try:
         callback(message)
-    except Exception as exc:
+    except BaseException as exc:
         # Note: the likelihood of this failing is extremely low. This just adds
         # a message to a queue, so if this doesn't work the world is in an
         # unrecoverable state and this thread should just bail.
diff --git a/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
@@ -77,8 +77,14 @@ def test__wrap_callback_errors_no_error():
     on_callback_error.assert_not_called()
 
 
-def test__wrap_callback_errors_error():
-    callback_error = ValueError("meep")
+@pytest.mark.parametrize(
+    "callback_error",
+    [
+        (ValueError("ValueError")),
+        (BaseException("BaseException")),
+    ],
+)
+def test__wrap_callback_errors_error(callback_error):
 
     msg = mock.create_autospec(message.Message, instance=True)
     callback = mock.Mock(side_effect=callback_error)
