diff --git a/google/api_core/bidi.py b/google/api_core/bidi.py index b002b409..bed4c70e 100644 --- a/google/api_core/bidi.py +++ b/google/api_core/bidi.py @@ -624,12 +624,15 @@ def on_response(response): ``open()``ed yet. on_response (Callable[[protobuf.Message], None]): The callback to be called for every response on the stream. + on_fatal_exception (Callable[[Exception], None]): The callback to + be called on fatal errors during consumption. Default None. """ - def __init__(self, bidi_rpc, on_response): + def __init__(self, bidi_rpc, on_response, on_fatal_exception=None): self._bidi_rpc = bidi_rpc self._on_response = on_response self._paused = False + self._on_fatal_exception = on_fatal_exception self._wake = threading.Condition() self._thread = None self._operational_lock = threading.Lock() @@ -676,6 +679,8 @@ def _thread_main(self, ready): exc, exc_info=True, ) + if self._on_fatal_exception is not None: + self._on_fatal_exception(exc) except Exception as exc: _LOGGER.exception( @@ -683,6 +688,8 @@ def _thread_main(self, ready): _BIDIRECTIONAL_CONSUMER_NAME, exc, ) + if self._on_fatal_exception is not None: + self._on_fatal_exception(exc) _LOGGER.info("%s exiting", _BIDIRECTIONAL_CONSUMER_NAME) @@ -694,8 +701,8 @@ def start(self): name=_BIDIRECTIONAL_CONSUMER_NAME, target=self._thread_main, args=(ready,), + daemon=True, ) - thread.daemon = True thread.start() # Other parts of the code rely on `thread.is_alive` which # isn't sufficient to know if a thread is active, just that it may @@ -706,7 +713,11 @@ def start(self): _LOGGER.debug("Started helper thread %s", thread.name) def stop(self): - """Stop consuming the stream and shutdown the background thread.""" + """Stop consuming the stream and shutdown the background thread. + + NOTE: Cannot be called within `_thread_main`, since it is not + possible to join a thread to itself. + """ with self._operational_lock: self._bidi_rpc.close() @@ -721,6 +732,7 @@ def stop(self): self._thread = None self._on_response = None + self._on_fatal_exception = None @property def is_active(self): diff --git a/tests/unit/test_bidi.py b/tests/unit/test_bidi.py index 0e7b018c..7640367c 100644 --- a/tests/unit/test_bidi.py +++ b/tests/unit/test_bidi.py @@ -918,3 +918,33 @@ def test_stop_error_logs(self, caplog): error_logs = [r.message for r in caplog.records if r.levelname == "ERROR"] assert not error_logs, f"Found unexpected ERROR logs: {error_logs}" bidi_rpc.is_active = False + + def test_fatal_exceptions_can_inform_consumer(self, caplog): + """ + https://github.com/googleapis/python-api-core/issues/820 + Exceptions thrown in the BackgroundConsumer not caught by `should_recover` / `should_terminate` + on the RPC should be bubbled back to the caller through `on_fatal_exception`, if passed. + """ + caplog.set_level(logging.DEBUG) + + for fatal_exception in ( + ValueError("some non-api error"), + exceptions.PermissionDenied("some api error"), + ): + bidi_rpc = mock.create_autospec(bidi.ResumableBidiRpc, instance=True) + bidi_rpc.is_active = True + on_response = mock.Mock(spec=["__call__"]) + + on_fatal_exception = mock.Mock(spec=["__call__"]) + + bidi_rpc.open.side_effect = fatal_exception + + consumer = bidi.BackgroundConsumer( + bidi_rpc, on_response, on_fatal_exception + ) + + consumer.start() + # let the background thread run for a while before exiting + time.sleep(0.1) + + on_fatal_exception.assert_called_once_with(fatal_exception)