refactor: export CertRotationInterceptor natively in google-auth

Author: agrawalradhika-cell

packages/google-auth/google/auth/transport/grpc.py

@@ -324,9 +324,9 @@ def my_client_cert_callback():
             "_is_recreation": True,  # Hidden flag to stop recursion
             **kwargs,
         }
-        interceptor = _MTLSCallInterceptor()
+        interceptor = CertRotationInterceptor()
 
-        wrapper = _MTLSRefreshingChannel(target, factory_args, channel, cached_cert)
+        wrapper = MTLSRefreshingChannel(target, factory_args, channel, cached_cert)
 
         interceptor._wrapper = wrapper
         return grpc.intercept_channel(wrapper, interceptor)
@@ -401,7 +401,7 @@ def is_mtls(self):
         return self._is_mtls
 
 
-class _MTLSCallInterceptor(
+class CertRotationInterceptor(
     grpc.UnaryUnaryClientInterceptor,
     grpc.UnaryStreamClientInterceptor,
     grpc.StreamUnaryClientInterceptor,
@@ -413,7 +413,7 @@ class _MTLSCallInterceptor(
     futures or iterators. Its primary role is to monitor responses for `UNAUTHENTICATED`
     errors. When an authentication failure occurs, it uses `_should_retry()` to check
     if a new mTLS certificate is available. If a new certificate is found, it signals
-    its associated `_MTLSRefreshingChannel` wrapper to refresh the underlying gRPC
+    its associated `MTLSRefreshingChannel` wrapper to refresh the underlying gRPC
     channel's credentials and automatically replays the failed RPC.
     """
 
@@ -485,7 +485,7 @@ def intercept_stream_stream(
         )
 
 
-class _MTLSRefreshingChannel(grpc.Channel):
+class MTLSRefreshingChannel(grpc.Channel):
     def __init__(self, target, factory_args, initial_channel, initial_cert):
         self._target = target
         self._factory_args = factory_args
@@ -681,6 +681,7 @@ def __init__(
         self._completion_event = threading.Event()
         self._done_callbacks = []
         self._terminal_exception = None
+        self._attempt_cert = None
 
         self._start_call()
 
@@ -734,7 +735,18 @@ def _on_inner_future_done(self, inner_future):
                     _LOGGER.warning("Callback failed: %s", e)
             return
 
-        exc = inner_future.exception()
+        try:
+            exc = inner_future.exception()
+        except Exception: # gracefully handle CancelledError or other future errors
+            with self._lock:
+                self._completion_event.set()
+                callbacks_to_fire = list(self._done_callbacks)
+            for fn in callbacks_to_fire:
+                try:
+                    fn(self)
+                except Exception as e:
+                    _LOGGER.warning("Callback failed: %s", e)
+            return
         if isinstance(exc, grpc.RpcError):
             status_code = exc.code()
 
@@ -745,10 +757,10 @@ def _on_inner_future_done(self, inner_future):
             )
 
             should_retry, call_cert, call_key, pwd = self._interceptor._should_retry(
-                status_code, self._retry_count, getattr(self, "_attempt_cert", None)
+                status_code, self._retry_count, self._attempt_cert
             )
             if can_replay and should_retry:
-                if getattr(self._interceptor, "_wrapper", None):
+                if self._interceptor._wrapper:
                     try:
                         self._interceptor._wrapper.refresh_logic(
                             1, call_cert, call_key, pwd
@@ -922,13 +934,14 @@ def __init__(
 
         self._is_completed = False
         self._done_callbacks = []
+        self._attempt_cert = None
 
         self._start_call()
 
     def _start_call(self):
         self._attempt_cert = (
             self._interceptor._wrapper._cached_cert
-            if getattr(self._interceptor, "_wrapper", None)
+            if self._interceptor._wrapper
             else None
         )
         with self._lock:
@@ -1022,12 +1035,12 @@ def __next__(self):
                 ) = self._interceptor._should_retry(
                     status_code,
                     self._retry_count,
-                    getattr(self, "_attempt_cert", None),
+                    self._attempt_cert,
                 )
 
                 if not self._yielded_any_response and can_replay and should_retry:
                     try:
-                        if getattr(self._interceptor, "_wrapper", None):
+                        if self._interceptor._wrapper:
                             self._interceptor._wrapper.refresh_logic(
                                 1, call_cert, call_key, pwd
                             )
@@ -1043,14 +1056,14 @@ def __next__(self):
                     continue
                 else:
                     # Non-retryable error, check if another rotation happened while we were finishing
-                    if getattr(self._interceptor, "_wrapper", None):
+                    if self._interceptor._wrapper:
                         (
                             chk_should_retry,
                             chk_cert,
                             chk_key,
                             chk_pwd,
                         ) = self._interceptor._should_retry(
-                            status_code, 0, getattr(self, "_attempt_cert", None)
+                            status_code, 0, self._attempt_cert
                         )
                         if chk_should_retry:
                             try:
@@ -1065,7 +1078,7 @@ def __next__(self):
 
     def add_done_callback(self, fn):
         with self._lock:
-            if getattr(self, "_is_completed", False):
+            if self._is_completed:
                 fire_now = True
             else:
                 self._done_callbacks.append(fn)

packages/google-auth/tests/transport/test_grpc.py

@@ -696,7 +696,7 @@ def test_get_client_ssl_credentials_auto_enablement(
 def test_interceptor_uses_factory_if_callable(mock_replayable):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
 
     call_no_factory = transport_grpc._RetryableStreamResponseIterator(
         continuation=mock.Mock(),
@@ -722,11 +722,11 @@ def generator_factory():
     assert call_factory._payload is None
 
 
-@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+@mock.patch("google.auth.transport.grpc.CertRotationInterceptor._should_retry")
 def test_factory_infinite_replay_on_error(mock_should_retry):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"
     mock_should_retry.side_effect = [
@@ -782,7 +782,7 @@ def test_refresh_logic_closes_old_channel(
 
     subscriber = mock.Mock()
 
-    refreshing_channel = transport_grpc._MTLSRefreshingChannel(
+    refreshing_channel = transport_grpc.MTLSRefreshingChannel(
         target="example.com:443",
         factory_args={},
         initial_channel=old_channel,
@@ -799,11 +799,11 @@ def test_refresh_logic_closes_old_channel(
     # old_channel.close.assert_called_once()  # Removed in PR 18019
 
 
-@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+@mock.patch("google.auth.transport.grpc.CertRotationInterceptor._should_retry")
 def test_unary_response_future_deadline_exceeded_on_retry(mock_should_retry):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"
     mock_should_retry.return_value = (True, b"cert", b"key", None)
@@ -847,11 +847,11 @@ def callback(f):
         future.result(timeout=1)
 
 
-@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+@mock.patch("google.auth.transport.grpc.CertRotationInterceptor._should_retry")
 def test_unary_response_future_cancelled(mock_should_retry):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"
 
@@ -886,11 +886,11 @@ def failing_callback(f):
     assert len(callbacks_fired) == 1
 
 
-@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+@mock.patch("google.auth.transport.grpc.CertRotationInterceptor._should_retry")
 def test_unary_response_future_rpc_error_retry_start_call_exception(mock_should_retry):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"
 
@@ -943,7 +943,7 @@ def test_start_call_wrapper_none():
     import pytest
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     if hasattr(interceptor, "_wrapper"):
         del interceptor._wrapper
 
@@ -963,7 +963,7 @@ def test_start_call_wrapper_none():
 def test_start_call_wrapper_none_branch():
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = None
 
     inner_future = mock.Mock()
@@ -980,11 +980,11 @@ def test_start_call_wrapper_none_branch():
     assert getattr(future, "_attempt_cert", "NOT_SET") is None
 
 
-@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+@mock.patch("google.auth.transport.grpc.CertRotationInterceptor._should_retry")
 def test_unary_response_future_rpc_error_no_wrapper(mock_should_retry):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = None
 
     mock_err = transport_grpc.grpc.RpcError()
@@ -1010,11 +1010,11 @@ def test_unary_response_future_rpc_error_no_wrapper(mock_should_retry):
         future._on_inner_future_done(inner_future)
 
 
-@mock.patch("google.auth.transport.grpc._MTLSCallInterceptor._should_retry")
+@mock.patch("google.auth.transport.grpc.CertRotationInterceptor._should_retry")
 def test_unary_response_future_rpc_error_should_not_retry(mock_should_retry):
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"
 
@@ -1047,7 +1047,7 @@ def test_mtls_call_interceptor_should_retry_cases():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
 
     assert interceptor._should_retry(grpc.StatusCode.UNAUTHENTICATED, 0, "cert1") == (
         False,
@@ -1102,7 +1102,7 @@ def test_mtls_call_interceptor_interceptors_methods():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
 
     def dummy_continuation(*args, **kwargs):
         return mock.Mock()
@@ -1135,7 +1135,7 @@ def test_mtls_refreshing_channel_refresh_logic_cases():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     with mock.patch(
         "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response"
     ) as mock_check, mock.patch("google.auth.transport.grpc.secure_authorized_channel"):
@@ -1185,7 +1185,7 @@ def test_mtls_refreshing_channel_subscribe_unsubscribe_close():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     cb = mock.Mock()
     channel.subscribe(cb)
     assert cb in channel._subscribers
@@ -1200,14 +1200,14 @@ def test_mtls_refreshing_channel_unary_unary():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     res = channel.unary_unary("method")
     assert res is not None
 
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     res = channel.unary_unary("method")
     assert res is not None
 
@@ -1216,7 +1216,7 @@ def test_mtls_refreshing_channel_unary_stream():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     res = channel.unary_stream("method")
     assert res is not None
 
@@ -1225,7 +1225,7 @@ def test_mtls_refreshing_channel_stream_unary():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     res = channel.stream_unary("method")
     assert res is not None
 
@@ -1234,7 +1234,7 @@ def test_mtls_refreshing_channel_stream_stream():
     from unittest import mock
     import google.auth.transport.grpc as transport_grpc
 
-    channel = transport_grpc._MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
+    channel = transport_grpc.MTLSRefreshingChannel("target", {}, mock.Mock(), b"cert1")
     res = channel.stream_stream("method")
     assert res is not None
 
@@ -1248,7 +1248,7 @@ def test_retryable_unary_response_future_methods():
     def dummy_continuation(*args, **kwargs):
         return mock_future
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"
 
@@ -1285,7 +1285,7 @@ def test_retryable_stream_response_iterator_methods():
     def dummy_continuation(*args, **kwargs):
         return mock_iterator
 
-    interceptor = transport_grpc._MTLSCallInterceptor()
+    interceptor = transport_grpc.CertRotationInterceptor()
     interceptor._wrapper = mock.Mock()
     interceptor._wrapper._cached_cert = "cert"