Add OTel auto-instrumentation (#163)

Signed-off-by: Anuraag Agrawal <anuraaga@gmail.com>

Author: anuraaga

connectrpc-otel/README.md

@@ -1,9 +1,7 @@
 # connectrpc-otel
 
-OpenTelemetry middleware for connect-python to generate server and client spans
-for ConnectRPC requests.
-
-Auto-instrumentation is currently not supported.
+OpenTelemetry instrumentation for connect-python to generate server and client spans and metrics
+for ConnectRPC requests with support for auto-instrumentation.
 
 ## Example
 

connectrpc-otel/connectrpc_otel/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
-__all__ = ["OpenTelemetryInterceptor"]
+__all__ = ["ConnectInstrumentor", "OpenTelemetryInterceptor"]
 
+from ._instrumentor import ConnectInstrumentor
 from ._interceptor import OpenTelemetryInterceptor

connectrpc-otel/connectrpc_otel/_instrumentor.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import sys
+from functools import partial
+from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from wrapt import register_post_import_hook, wrap_function_wrapper
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Collection
+    from types import ModuleType
+
+    from opentelemetry.metrics import MeterProvider
+    from opentelemetry.trace import TracerProvider
+
+    from connectrpc.interceptor import Interceptor, InterceptorSync
+
+_instruments = ("connect-python>=0.9.0",)
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class ConnectInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: object) -> None:
+        self._meter_provider = cast(
+            "MeterProvider | None", kwargs.get("meter_provider")
+        )
+        self._tracer_provider = cast(
+            "TracerProvider | None", kwargs.get("tracer_provider")
+        )
+
+        register_post_import_hook(self._patch_client, "connectrpc.client")
+        register_post_import_hook(self._patch_server, "connectrpc.server")
+
+    def _uninstrument(self, **kwargs: object) -> None:
+        # TODO: Remove sys.modules check after
+        # https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4321
+        if "connectrpc.client" in sys.modules:
+            unwrap("connectrpc.client.ConnectClient", "__init__")
+            unwrap("connectrpc.client.ConnectClientSync", "__init__")
+        if "connectrpc.server" in sys.modules:
+            unwrap("connectrpc.server.ConnectASGIApplication", "__init__")
+            unwrap("connectrpc.server.ConnectWSGIApplication", "__init__")
+
+    def _patch_client(self, module: ModuleType) -> None:
+        wrap_function_wrapper(
+            module, "ConnectClient.__init__", partial(self._init_wrapper, client=True)
+        )
+        wrap_function_wrapper(
+            module,
+            "ConnectClientSync.__init__",
+            partial(self._init_wrapper, client=True),
+        )
+
+    def _patch_server(self, module: ModuleType) -> None:
+        wrap_function_wrapper(
+            module,
+            "ConnectASGIApplication.__init__",
+            partial(self._init_wrapper, client=False),
+        )
+        wrap_function_wrapper(
+            module,
+            "ConnectWSGIApplication.__init__",
+            partial(self._init_wrapper, client=False),
+        )
+
+    def _init_wrapper(
+        self,
+        wrapped: Callable[P, R],
+        _instance: object,
+        args: tuple,
+        kwargs: dict,
+        *,
+        client: bool,
+    ) -> R:
+        # Instrumentation doesn't eager import the library it instruments
+        from connectrpc_otel import OpenTelemetryInterceptor  # noqa: PLC0415
+
+        interceptors: list[Interceptor | InterceptorSync] | None = kwargs.get(
+            "interceptors"
+        )
+        interceptors = [] if interceptors is None else [*interceptors]
+        if any(isinstance(i, OpenTelemetryInterceptor) for i in interceptors):
+            return wrapped(*args, **kwargs)
+
+        kwargs["interceptors"] = interceptors
+        # OpenTelemetry interceptor should be first so i.e. logging interceptors
+        # have trace IDs.
+        interceptors.insert(
+            0,
+            OpenTelemetryInterceptor(
+                tracer_provider=self._tracer_provider,
+                meter_provider=self._meter_provider,
+                client=client,
+            ),
+        )
+
+        return wrapped(*args, **kwargs)

connectrpc-otel/pyproject.toml

@@ -35,10 +35,23 @@ classifiers = [
     "Programming Language :: Python :: 3.14",
     "Topic :: System :: Networking",
 ]
-dependencies = ["connect-python>=0.8.0", "opentelemetry-api>=1.39.1"]
+dependencies = [
+    "opentelemetry-api>=1.39.1",
+    "opentelemetry-instrumentation==0.60b1",
+]
 
 [dependency-groups]
 dev = [
+    # auto-instrumentation shouldn't include a dependency on the library being instrumented
+    # so it can be packaged in auto-instrumentation distros for general use.
+    #
+    # This library also supports manually initializing the interceptor.
+    # For users that manually instrument by initializing the interceptor, it is easier to
+    # have a real dependency to avoid any possible version conflicts. But for Python,
+    # the ecosystem vastly favors auto-instrumentation, and it is easier to add than remove
+    # a transitive dependency, so we go ahead and leave it out.
+    "connect-python>=0.8.0",
+
     "opentelemetry-sdk==1.39.1",
     "opentelemetry-instrumentation-asgi==0.60b1",
     "opentelemetry-instrumentation-wsgi==0.60b1",
@@ -47,6 +60,9 @@ dev = [
     "pytest",
 ]
 
+[project.entry-points.opentelemetry_instrumentor]
+connectrpc = "connectrpc_otel:ConnectInstrumentor"
+
 [project.urls]
 Homepage = "https://github.com/connectrpc/connect-python"
 Repository = "https://github.com/connectrpc/connect-python"

connectrpc-otel/test/__init__.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import pytest
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import InMemoryMetricReader
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+
+@pytest.fixture
+def span_exporter() -> InMemorySpanExporter:
+    return InMemorySpanExporter()
+
+
+@pytest.fixture
+def tracer_provider(span_exporter: InMemorySpanExporter) -> TracerProvider:
+    tp = TracerProvider()
+    tp.add_span_processor(SimpleSpanProcessor(span_exporter))
+    return tp
+
+
+@pytest.fixture
+def metric_reader() -> InMemoryMetricReader:
+    return InMemoryMetricReader()
+
+
+@pytest.fixture
+def meter_provider(metric_reader: InMemoryMetricReader) -> MeterProvider:
+    return MeterProvider(metric_readers=[metric_reader])