diff --git a/packages/google-api-core/google/api_core/_observability.py b/packages/google-api-core/google/api_core/_observability.py index b1b71b056658..701829796827 100644 --- a/packages/google-api-core/google/api_core/_observability.py +++ b/packages/google-api-core/google/api_core/_observability.py @@ -16,16 +16,22 @@ """OpenTelemetry helpers for resolving and instantiating interceptors.""" -from typing import Any, Optional +import functools +from typing import TYPE_CHECKING, Any, Callable from google.api_core import _feature_gating_helpers from google.api_core.client_options import ClientOptions +if TYPE_CHECKING: + from google.api_core.grpc_helpers import ChannelWrapperCallable +else: + ChannelWrapperCallable = Callable[[Any], Any] + _TRACER_PROVIDER = "tracer_provider" def is_otel_capabilities_enabled( - client_options: Optional[ClientOptions | dict[str, Any]] = None, + client_options: ClientOptions | dict[str, Any] | None = None, env_var: str = "GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", ) -> bool: """Checks if OTel capabilities are enabled and installed. @@ -54,26 +60,19 @@ def is_otel_capabilities_enabled( return False -def apply_otel_capabilities_to_channel( - channel: Any, - client_options: Optional[ClientOptions | dict[str, Any]] = None, +def _get_otel_interceptor( + client_options: ClientOptions | dict[str, Any] | None = None, + is_async: bool = False, ) -> Any: - """Applies OTel capabilities (like tracing) to the channel. - - Precondition: This function assumes `is_otel_capabilities_enabled` has already - been called and returned `True`, i.e. in the Client. At this time - this function is not intended to be standalone. + """Instantiates a sync or async OpenTelemetry gRPC client interceptor. Args: - channel: The raw gRPC channel to wrap. client_options: The client options object or dictionary. + is_async: If True, returns an async interceptor (`aio_client_interceptor`), + otherwise returns a sync interceptor (`client_interceptor`). Returns: - Any: The intercepted channel. - - Raises: - ImportError: If OpenTelemetry packages are not installed and this function - is called directly (bypassing the precondition). + Any: The instantiated OpenTelemetry client interceptor. """ import opentelemetry.instrumentation.grpc as otel_grpc # type: ignore[import-not-found] @@ -83,7 +82,47 @@ def apply_otel_capabilities_to_channel( elif client_options is not None: tracer_provider = getattr(client_options, _TRACER_PROVIDER, None) - interceptor = otel_grpc.client_interceptor(tracer_provider=tracer_provider) + if is_async: + return otel_grpc.aio_client_interceptor(tracer_provider=tracer_provider) + return otel_grpc.client_interceptor(tracer_provider=tracer_provider) + + +def get_otel_channel_wrapper( + client_options: ClientOptions | dict[str, Any] | None = None, +) -> ChannelWrapperCallable | None: + """Returns a channel wrapper callable that wraps a sync gRPC channel with OpenTelemetry tracing. + + Args: + client_options: The client options object or dictionary used for feature gating + and extracting the tracer provider. + + Returns: + Optional[ChannelWrapperCallable]: A channel-wrapping callable if OpenTelemetry + tracing is enabled and installed, None otherwise. + """ + if not is_otel_capabilities_enabled(client_options): + return None + + import opentelemetry.instrumentation.grpc as otel_grpc # type: ignore[import-not-found] + + interceptor = _get_otel_interceptor(client_options, is_async=False) + return functools.partial(otel_grpc.intercept_channel, interceptor=interceptor) + + +def get_otel_async_interceptor( + client_options: ClientOptions | dict[str, Any] | None = None, +) -> Any | None: + """Returns an async gRPC client interceptor for OpenTelemetry tracing. + + Args: + client_options: The client options object or dictionary used for feature gating + and extracting the tracer provider. + + Returns: + Optional[Any]: An instantiated OpenTelemetry async client interceptor + if tracing is enabled and installed, None otherwise. + """ + if not is_otel_capabilities_enabled(client_options): + return None - # We use OTel's own compatible applier to avoid standard gRPC TypeError. - return otel_grpc.intercept_channel(channel, interceptor) + return _get_otel_interceptor(client_options, is_async=True) diff --git a/packages/google-api-core/tests/unit/test_observability.py b/packages/google-api-core/tests/unit/test_observability.py index f0ebe0afc14d..4005aa14fc8e 100644 --- a/packages/google-api-core/tests/unit/test_observability.py +++ b/packages/google-api-core/tests/unit/test_observability.py @@ -87,16 +87,57 @@ def test_is_otel_capabilities_enabled_experimental_enabled_with_config(monkeypat assert _observability.is_otel_capabilities_enabled(options) -def test_apply_otel_capabilities_to_channel_enabled_otel_installed(monkeypatch): - mock_channel = mock.Mock() - mock_intercepted_channel = mock.Mock() +def test_get_otel_interceptor_sync_default(monkeypatch): + mock_otel = mock.Mock() + mock_otel_grpc = mock_otel.instrumentation.grpc + mock_interceptor = mock.Mock() + mock_otel_grpc.client_interceptor.return_value = mock_interceptor + + monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation", mock_otel.instrumentation + ) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc + ) + + result = _observability._get_otel_interceptor() + assert result is mock_interceptor + mock_otel_grpc.client_interceptor.assert_called_once_with(tracer_provider=None) + + +def test_get_otel_interceptor_sync_config(monkeypatch): + mock_tracer_provider = object() + options = ClientOptions(tracer_provider=mock_tracer_provider) mock_otel = mock.Mock() mock_otel_grpc = mock_otel.instrumentation.grpc mock_interceptor = mock.Mock() + mock_otel_grpc.client_interceptor.return_value = mock_interceptor + + monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation", mock_otel.instrumentation + ) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc + ) + result = _observability._get_otel_interceptor(client_options=options) + assert result is mock_interceptor + mock_otel_grpc.client_interceptor.assert_called_once_with( + tracer_provider=mock_tracer_provider + ) + + +def test_get_otel_interceptor_sync_dict_config(monkeypatch): + mock_tracer_provider = object() + options = {"tracer_provider": mock_tracer_provider} + + mock_otel = mock.Mock() + mock_otel_grpc = mock_otel.instrumentation.grpc + mock_interceptor = mock.Mock() mock_otel_grpc.client_interceptor.return_value = mock_interceptor - mock_otel_grpc.intercept_channel.return_value = mock_intercepted_channel monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) monkeypatch.setitem( @@ -106,29 +147,62 @@ def test_apply_otel_capabilities_to_channel_enabled_otel_installed(monkeypatch): sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc ) - result = _observability.apply_otel_capabilities_to_channel(mock_channel) + result = _observability._get_otel_interceptor(client_options=options) + assert result is mock_interceptor + mock_otel_grpc.client_interceptor.assert_called_once_with( + tracer_provider=mock_tracer_provider + ) + - assert result is mock_intercepted_channel - mock_otel_grpc.client_interceptor.assert_called_once_with(tracer_provider=None) - mock_otel_grpc.intercept_channel.assert_called_once_with( - mock_channel, mock_interceptor +def test_get_otel_interceptor_async(monkeypatch): + mock_tracer_provider = object() + options = ClientOptions(tracer_provider=mock_tracer_provider) + + mock_otel = mock.Mock() + mock_otel_grpc = mock_otel.instrumentation.grpc + mock_async_interceptor = mock.Mock() + mock_otel_grpc.aio_client_interceptor.return_value = mock_async_interceptor + + monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation", mock_otel.instrumentation ) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc + ) + + result = _observability._get_otel_interceptor(client_options=options, is_async=True) + assert result is mock_async_interceptor + mock_otel_grpc.aio_client_interceptor.assert_called_once_with( + tracer_provider=mock_tracer_provider + ) + + +def test_get_otel_channel_wrapper_disabled(monkeypatch): + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "false") + assert _observability.get_otel_channel_wrapper() is None + + +def test_get_otel_channel_wrapper_otel_missing(monkeypatch): + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "true") + monkeypatch.setitem(sys.modules, "opentelemetry.instrumentation.grpc", None) + assert _observability.get_otel_channel_wrapper() is None -def test_apply_otel_capabilities_to_channel_enabled_via_config(monkeypatch): - # Tracing enabled via config (tracer_provider is set) +def test_get_otel_channel_wrapper_enabled(monkeypatch): + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "true") mock_tracer_provider = object() options = ClientOptions(tracer_provider=mock_tracer_provider) - mock_channel = mock.Mock() - mock_intercepted_channel = mock.Mock() + mock_raw_channel = mock.Mock(name="raw_channel") + mock_wrapped_channel = mock.Mock(name="wrapped_channel") mock_otel = mock.Mock() mock_otel_grpc = mock_otel.instrumentation.grpc - mock_interceptor = mock.Mock() + mock_interceptor = mock.Mock(name="otel_interceptor") mock_otel_grpc.client_interceptor.return_value = mock_interceptor - mock_otel_grpc.intercept_channel.return_value = mock_intercepted_channel + mock_otel_grpc.intercept_channel.return_value = mock_wrapped_channel monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) monkeypatch.setitem( @@ -138,33 +212,38 @@ def test_apply_otel_capabilities_to_channel_enabled_via_config(monkeypatch): sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc ) - result = _observability.apply_otel_capabilities_to_channel( - mock_channel, client_options=options - ) + wrapper = _observability.get_otel_channel_wrapper(client_options=options) + assert callable(wrapper) - assert result is mock_intercepted_channel mock_otel_grpc.client_interceptor.assert_called_once_with( tracer_provider=mock_tracer_provider ) + + result = wrapper(mock_raw_channel) + assert result is mock_wrapped_channel mock_otel_grpc.intercept_channel.assert_called_once_with( - mock_channel, mock_interceptor + mock_raw_channel, interceptor=mock_interceptor ) -def test_apply_otel_capabilities_to_channel_enabled_via_dict_config(monkeypatch): - # Tracing enabled via dict config +def test_get_otel_channel_wrapper_with_apply_channel_wrappers(monkeypatch): + """Proves that get_otel_channel_wrapper integrates seamlessly into apply_channel_wrappers.""" + pytest.importorskip("grpc") + from google.api_core import grpc_helpers + + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "true") mock_tracer_provider = object() - options = {"tracer_provider": mock_tracer_provider} + options = ClientOptions(tracer_provider=mock_tracer_provider) - mock_channel = mock.Mock() - mock_intercepted_channel = mock.Mock() + mock_raw_channel = mock.Mock(name="raw_channel") + mock_wrapped_channel = mock.Mock(name="wrapped_channel") mock_otel = mock.Mock() mock_otel_grpc = mock_otel.instrumentation.grpc - mock_interceptor = mock.Mock() + mock_interceptor = mock.Mock(name="otel_interceptor") mock_otel_grpc.client_interceptor.return_value = mock_interceptor - mock_otel_grpc.intercept_channel.return_value = mock_intercepted_channel + mock_otel_grpc.intercept_channel.return_value = mock_wrapped_channel monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) monkeypatch.setitem( @@ -174,14 +253,50 @@ def test_apply_otel_capabilities_to_channel_enabled_via_dict_config(monkeypatch) sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc ) - result = _observability.apply_otel_capabilities_to_channel( - mock_channel, client_options=options - ) + otel_wrapper = _observability.get_otel_channel_wrapper(client_options=options) + assert callable(otel_wrapper) - assert result is mock_intercepted_channel - mock_otel_grpc.client_interceptor.assert_called_once_with( - tracer_provider=mock_tracer_provider + result = grpc_helpers.apply_channel_wrappers( + mock_raw_channel, wrappers=[otel_wrapper] ) + assert result is mock_wrapped_channel mock_otel_grpc.intercept_channel.assert_called_once_with( - mock_channel, mock_interceptor + mock_raw_channel, interceptor=mock_interceptor + ) + + +def test_get_otel_async_interceptor_disabled(monkeypatch): + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "false") + assert _observability.get_otel_async_interceptor() is None + + +def test_get_otel_async_interceptor_otel_missing(monkeypatch): + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "true") + monkeypatch.setitem(sys.modules, "opentelemetry.instrumentation.grpc", None) + assert _observability.get_otel_async_interceptor() is None + + +def test_get_otel_async_interceptor_enabled(monkeypatch): + monkeypatch.setenv("GOOGLE_SDK_EXPERIMENTAL_PYTHON_TRACING_ENABLED", "true") + mock_tracer_provider = object() + options = ClientOptions(tracer_provider=mock_tracer_provider) + + mock_async_interceptor = mock.Mock(name="otel_async_interceptor") + + mock_otel = mock.Mock() + mock_otel_grpc = mock_otel.instrumentation.grpc + mock_otel_grpc.aio_client_interceptor.return_value = mock_async_interceptor + + monkeypatch.setitem(sys.modules, "opentelemetry", mock_otel) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation", mock_otel.instrumentation + ) + monkeypatch.setitem( + sys.modules, "opentelemetry.instrumentation.grpc", mock_otel_grpc + ) + + result = _observability.get_otel_async_interceptor(client_options=options) + assert result is mock_async_interceptor + mock_otel_grpc.aio_client_interceptor.assert_called_once_with( + tracer_provider=mock_tracer_provider )