feat(api-core): add ChannelWrapper and apply_channel_wrappers helper

Author: chalmerlowe

packages/google-api-core/google/api_core/grpc_helpers.py

@@ -17,7 +17,16 @@
 import collections
 import functools
 import warnings
-from typing import Generic, Iterator, Optional, Sequence, TypeVar, Union
+from typing import (
+    Callable,
+    Generic,
+    Iterator,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    get_args,
+)
 
 import google.auth
 import google.auth.credentials
@@ -41,6 +50,15 @@
     grpc.StreamStreamClientInterceptor,
 ]
 
+# Runtime tuple of gRPC client interceptor base classes for isinstance checks
+_CLIENT_INTERCEPTOR_CLASSES = get_args(ClientInterceptor)
+
+# Type alias representing a channel-wrapping callable
+ChannelWrapperCallable = Callable[[grpc.Channel], grpc.Channel]
+
+# Generic type alias representing any channel wrapper (interceptor or callable)
+ChannelWrapper = Union[ClientInterceptor, ChannelWrapperCallable]
+
 
 def _patch_callable_name(callable_):
     """Fix-up gRPC callable attributes.
@@ -426,27 +444,44 @@ def _modify_target_for_direct_path(target: str) -> str:
     return target
 
 
-def apply_interceptors(
+def apply_channel_wrappers(
     channel: grpc.Channel,
-    interceptors: Optional[Sequence[ClientInterceptor]] = None,
+    wrappers: Optional[Sequence[ChannelWrapper]] = None,
 ) -> grpc.Channel:
-    """Applies client interceptors to a gRPC channel.
+    """Applies channel wrappers (client interceptors or channel-wrapping callables) to a gRPC channel.
 
-    The first interceptor in the sequence is the outermost layer: it
-    executes first on outbound requests and last on inbound responses.
+    Executes in reverse order so the first wrapper in the sequence becomes the
+    outermost layer on outbound requests and the innermost layer on inbound responses.
 
     Args:
-        channel (grpc.Channel): The channel to intercept.
-        interceptors (Optional[Sequence[ClientInterceptor]]): An optional sequence
-            of client interceptors to apply.
+        channel (grpc.Channel): The channel to wrap.
+        wrappers (Optional[Sequence[ChannelWrapper]]):
+            An optional sequence of client interceptors or channel-wrapping
+            callables to apply.
 
     Returns:
-        grpc.Channel: The intercepted channel, or the original channel if no
-            interceptors were provided.
+        grpc.Channel: The wrapped channel, or the original channel if no
+            wrappers were provided.
+
+    Raises:
+        TypeError: If an item in ``wrappers`` is neither a gRPC ClientInterceptor
+            nor a Callable[[Channel], Channel].
     """
-    if interceptors:
-        return grpc.intercept_channel(channel, *interceptors)
-    return channel
+    if not wrappers:
+        return channel
+
+    modified_channel = channel
+    for wrapper in reversed(list(wrappers)):
+        if isinstance(wrapper, _CLIENT_INTERCEPTOR_CLASSES):
+            modified_channel = grpc.intercept_channel(modified_channel, wrapper)
+        elif callable(wrapper):
+            modified_channel = wrapper(modified_channel)
+        else:
+            raise TypeError(
+                f"Expected ChannelWrapper (ClientInterceptor or Callable[[Channel], Channel]), got {type(wrapper).__name__}"
+            )
+
+    return modified_channel
 
 
 _MethodCall = collections.namedtuple(

packages/google-api-core/tests/unit/test_grpc_helpers.py

@@ -933,32 +933,100 @@ def test_close(self):
         assert channel.close() is None
 
 
-@pytest.mark.parametrize("falsy_interceptors", [None, [], ()])
-def test_apply_interceptors_passthrough(falsy_interceptors):
-    """Verify that falsy or empty interceptor sequences return the channel unmodified."""
+@pytest.mark.parametrize("falsy_wrappers", [None, [], ()])
+def test_apply_channel_wrappers_passthrough(falsy_wrappers):
+    """Verify that falsy or empty wrapper sequences return the channel unmodified."""
     mock_base_channel = mock.Mock(name="base_channel")
-    result = grpc_helpers.apply_interceptors(mock_base_channel, falsy_interceptors)
+    result = grpc_helpers.apply_channel_wrappers(mock_base_channel, falsy_wrappers)
     assert result is mock_base_channel
 
 
-@pytest.mark.parametrize("count", [1, 2, 3])
-def test_apply_interceptors_wrapping(count):
-    """Verify that interceptors are passed to grpc.intercept_channel unpacked in a single call.
+def test_apply_channel_wrappers_grpc_client_interceptors():
+    """Verify that standard gRPC ClientInterceptor instances are applied via grpc.intercept_channel."""
+
+    class DummyUnaryInterceptor(grpc.UnaryUnaryClientInterceptor):
+        def intercept_unary_unary(self, continuation, client_call_details, request):
+            return continuation(client_call_details, request)
+
+    class DummyStreamInterceptor(grpc.StreamStreamClientInterceptor):
+        def intercept_stream_stream(
+            self, continuation, client_call_details, request_iterator
+        ):
+            return continuation(client_call_details, request_iterator)
+
+    interceptor1 = DummyUnaryInterceptor()
+    interceptor2 = DummyStreamInterceptor()
 
-    When given a sequence of N interceptors [i_0, i_1, ..., i_{N-1}], apply_interceptors
-    must pass the base channel and all interceptors unpacked (*interceptors) to
-    grpc.intercept_channel.
-    """
     mock_base_channel = mock.Mock(name="base_channel")
-    mock_wrapped_channel = mock.Mock(name="wrapped_channel")
-    mock_interceptors = [mock.Mock(name=f"interceptor_{i}") for i in range(count)]
+    mock_chan_after_i2 = mock.Mock(name="chan_after_i2")
+    mock_chan_after_i1 = mock.Mock(name="chan_after_i1")
 
     with mock.patch(
-        "grpc.intercept_channel", return_value=mock_wrapped_channel
-    ) as mock_intercept_channel:
-        result = grpc_helpers.apply_interceptors(mock_base_channel, mock_interceptors)
+        "grpc.intercept_channel",
+        side_effect=[mock_chan_after_i2, mock_chan_after_i1],
+    ) as mock_intercept:
+        result = grpc_helpers.apply_channel_wrappers(
+            mock_base_channel, [interceptor1, interceptor2]
+        )
 
-    assert result is mock_wrapped_channel
-    mock_intercept_channel.assert_called_once_with(
-        mock_base_channel, *mock_interceptors
+    assert result is mock_chan_after_i1
+    assert mock_intercept.call_count == 2
+    # Executed in reverse order so interceptor1 is outermost
+    mock_intercept.assert_has_calls(
+        [
+            mock.call(mock_base_channel, interceptor2),
+            mock.call(mock_chan_after_i2, interceptor1),
+        ]
     )
+
+
+def test_apply_channel_wrappers_callables():
+    """Verify that channel-wrapping callables Callable[[Channel], Channel] are invoked in sequence."""
+    mock_base_channel = mock.Mock(name="base_channel")
+    mock_chan_1 = mock.Mock(name="chan_1")
+    mock_chan_2 = mock.Mock(name="chan_2")
+
+    wrapper1 = mock.Mock(side_effect=lambda ch: mock_chan_2)
+    wrapper2 = mock.Mock(side_effect=lambda ch: mock_chan_1)
+
+    result = grpc_helpers.apply_channel_wrappers(
+        mock_base_channel, [wrapper1, wrapper2]
+    )
+
+    assert result is mock_chan_2
+    # Executed in reverse order: wrapper2 runs first on base channel, then wrapper1
+    wrapper2.assert_called_once_with(mock_base_channel)
+    wrapper1.assert_called_once_with(mock_chan_1)
+
+
+def test_apply_channel_wrappers_interspersed():
+    """Verify that a mixed sequence of gRPC interceptors and channel wrapper callables are applied."""
+
+    class DummyUnaryInterceptor(grpc.UnaryUnaryClientInterceptor):
+        def intercept_unary_unary(self, continuation, client_call_details, request):
+            return continuation(client_call_details, request)
+
+    interceptor = DummyUnaryInterceptor()
+    mock_base_channel = mock.Mock(name="base_channel")
+    mock_chan_after_wrapper = mock.Mock(name="chan_after_wrapper")
+    mock_chan_after_interceptor = mock.Mock(name="chan_after_interceptor")
+
+    wrapper = mock.Mock(return_value=mock_chan_after_wrapper)
+
+    with mock.patch(
+        "grpc.intercept_channel", return_value=mock_chan_after_interceptor
+    ) as mock_intercept:
+        result = grpc_helpers.apply_channel_wrappers(
+            mock_base_channel, [interceptor, wrapper]
+        )
+
+    assert result is mock_chan_after_interceptor
+    wrapper.assert_called_once_with(mock_base_channel)
+    mock_intercept.assert_called_once_with(mock_chan_after_wrapper, interceptor)
+
+
+def test_apply_channel_wrappers_invalid_type_raises():
+    """Verify that passing an invalid object that is neither an interceptor nor callable raises TypeError."""
+    mock_base_channel = mock.Mock(name="base_channel")
+    with pytest.raises(TypeError, match="Expected ChannelWrapper"):
+        grpc_helpers.apply_channel_wrappers(mock_base_channel, [12345])