diff --git a/packages/google-api-core/google/api_core/grpc_helpers.py b/packages/google-api-core/google/api_core/grpc_helpers.py index 263079e7d1f7..2f0ef9631dd8 100644 --- a/packages/google-api-core/google/api_core/grpc_helpers.py +++ b/packages/google-api-core/google/api_core/grpc_helpers.py @@ -17,7 +17,7 @@ import collections import functools import warnings -from typing import Generic, Iterator, Optional, TypeVar +from typing import Generic, Iterator, Optional, Sequence, TypeVar, Union import google.auth import google.auth.credentials @@ -25,7 +25,6 @@ import google.auth.transport.requests import google.protobuf import grpc - from google.api_core import exceptions, general_helpers # The list of gRPC Callable interfaces that return iterators. @@ -34,6 +33,14 @@ # denotes the proto response type for grpc calls P = TypeVar("P") +# Type alias representing any client-side gRPC interceptor +ClientInterceptor = Union[ + grpc.UnaryUnaryClientInterceptor, + grpc.UnaryStreamClientInterceptor, + grpc.StreamUnaryClientInterceptor, + grpc.StreamStreamClientInterceptor, +] + def _patch_callable_name(callable_): """Fix-up gRPC callable attributes. @@ -419,6 +426,29 @@ def _modify_target_for_direct_path(target: str) -> str: return target +def apply_interceptors( + channel: grpc.Channel, + interceptors: Optional[Sequence[ClientInterceptor]] = None, +) -> grpc.Channel: + """Applies client interceptors 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. + + Args: + channel (grpc.Channel): The channel to intercept. + interceptors (Optional[Sequence[ClientInterceptor]]): An optional sequence + of client interceptors to apply. + + Returns: + grpc.Channel: The intercepted channel, or the original channel if no + interceptors were provided. + """ + if interceptors: + return grpc.intercept_channel(channel, *interceptors) + return channel + + _MethodCall = collections.namedtuple( "_MethodCall", ("request", "timeout", "metadata", "credentials", "compression") ) diff --git a/packages/google-api-core/tests/unit/test_grpc_helpers.py b/packages/google-api-core/tests/unit/test_grpc_helpers.py index 69281d58109b..9e41bab853df 100644 --- a/packages/google-api-core/tests/unit/test_grpc_helpers.py +++ b/packages/google-api-core/tests/unit/test_grpc_helpers.py @@ -24,9 +24,8 @@ pytest.skip("No GRPC", allow_module_level=True) import google.auth.credentials -from google.longrunning import operations_pb2 - from google.api_core import exceptions, grpc_helpers +from google.longrunning import operations_pb2 def test__patch_callable_name(): @@ -932,3 +931,34 @@ def test_subscribe_unsubscribe(self): def test_close(self): channel = grpc_helpers.ChannelStub() 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.""" + mock_base_channel = mock.Mock(name="base_channel") + result = grpc_helpers.apply_interceptors(mock_base_channel, falsy_interceptors) + 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. + + 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)] + + 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) + + assert result is mock_wrapped_channel + mock_intercept_channel.assert_called_once_with( + mock_base_channel, *mock_interceptors + )