[opentelemetry-instrumentation-httpx] fix mixing of sync and non asyn…

…c hooks (#1920) * fix: sync response hooks being used on httpx.AsyncClient * docs: add changelog * docs: improved docs a bit more * docs: fix variable name --------- Co-authored-by: Diego Hurtado <[email protected]> Co-authored-by: Shalev Roda <[email protected]>
open-telemetry · Nov 21, 2023 · 1b1c38d · 1b1c38d
1 parent 773e431
commit 1b1c38d
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `opentelemetry-instrumentation` Added Otel semantic convention opt-in mechanism
   ([#1987](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1987))
+- `opentelemetry-instrumentation-httpx` Fix mixing async and non async hooks
+  ([#1920](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1920))
 
 ### Fixed
 

diff --git a/instrumentation/opentelemetry-instrumentation-httpx/README.rst b/instrumentation/opentelemetry-instrumentation-httpx/README.rst
@@ -136,15 +136,29 @@ The hooks can be configured as follows:
         # status_code, headers, stream, extensions = response
         pass
 
-    HTTPXClientInstrumentor().instrument(request_hook=request_hook, response_hook=response_hook)
+    async def async_request_hook(span, request):
+        # method, url, headers, stream, extensions = request
+        pass
+
+    async def async_response_hook(span, request, response):
+        # method, url, headers, stream, extensions = request
+        # status_code, headers, stream, extensions = response
+        pass
+
+    HTTPXClientInstrumentor().instrument(
+        request_hook=request_hook,
+        response_hook=response_hook,
+        async_request_hook=async_request_hook,
+        async_response_hook=async_response_hook
+    )
 
 
 Or if you are using the transport classes directly:
 
 
 .. code-block:: python
 
-    from opentelemetry.instrumentation.httpx import SyncOpenTelemetryTransport
+    from opentelemetry.instrumentation.httpx import SyncOpenTelemetryTransport, AsyncOpenTelemetryTransport
 
     def request_hook(span, request):
         # method, url, headers, stream, extensions = request
@@ -155,13 +169,29 @@ Or if you are using the transport classes directly:
         # status_code, headers, stream, extensions = response
         pass
 
+    async def async_request_hook(span, request):
+        # method, url, headers, stream, extensions = request
+        pass
+
+    async def async_response_hook(span, request, response):
+        # method, url, headers, stream, extensions = request
+        # status_code, headers, stream, extensions = response
+        pass
+
     transport = httpx.HTTPTransport()
     telemetry_transport = SyncOpenTelemetryTransport(
         transport,
         request_hook=request_hook,
         response_hook=response_hook
     )
 
+    async_transport = httpx.AsyncHTTPTransport()
+    async_telemetry_transport = AsyncOpenTelemetryTransport(
+        async_transport,
+        request_hook=async_request_hook,
+        response_hook=async_response_hook
+    )
+
 
 References
 ----------

diff --git a/...n/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py b/...n/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py
@@ -131,15 +131,29 @@ def response_hook(span, request, response):
         # status_code, headers, stream, extensions = response
         pass
 
-    HTTPXClientInstrumentor().instrument(request_hook=request_hook, response_hook=response_hook)
+    async def async_request_hook(span, request):
+        # method, url, headers, stream, extensions = request
+        pass
+
+    async def async_response_hook(span, request, response):
+        # method, url, headers, stream, extensions = request
+        # status_code, headers, stream, extensions = response
+        pass
+
+    HTTPXClientInstrumentor().instrument(
+        request_hook=request_hook,
+        response_hook=response_hook,
+        async_request_hook=async_request_hook,
+        async_response_hook=async_response_hook
+    )
 
 
 Or if you are using the transport classes directly:
 
 
 .. code-block:: python
 
-    from opentelemetry.instrumentation.httpx import SyncOpenTelemetryTransport
+    from opentelemetry.instrumentation.httpx import SyncOpenTelemetryTransport, AsyncOpenTelemetryTransport
 
     def request_hook(span, request):
         # method, url, headers, stream, extensions = request
@@ -150,13 +164,29 @@ def response_hook(span, request, response):
         # status_code, headers, stream, extensions = response
         pass
 
+    async def async_request_hook(span, request):
+        # method, url, headers, stream, extensions = request
+        pass
+
+    async def async_response_hook(span, request, response):
+        # method, url, headers, stream, extensions = request
+        # status_code, headers, stream, extensions = response
+        pass
+
     transport = httpx.HTTPTransport()
     telemetry_transport = SyncOpenTelemetryTransport(
         transport,
         request_hook=request_hook,
         response_hook=response_hook
     )
 
+    async_transport = httpx.AsyncHTTPTransport()
+    async_telemetry_transport = AsyncOpenTelemetryTransport(
+        async_transport,
+        request_hook=async_request_hook,
+        response_hook=async_response_hook
+    )
+
 API
 ---
 """
@@ -377,8 +407,8 @@ def __init__(
         self,
         transport: httpx.AsyncBaseTransport,
         tracer_provider: typing.Optional[TracerProvider] = None,
-        request_hook: typing.Optional[RequestHook] = None,
-        response_hook: typing.Optional[ResponseHook] = None,
+        request_hook: typing.Optional[AsyncRequestHook] = None,
+        response_hook: typing.Optional[AsyncResponseHook] = None,
     ):
         self._transport = transport
         self._tracer = get_tracer(
@@ -511,21 +541,27 @@ def _instrument(self, **kwargs):
         Args:
             **kwargs: Optional arguments
                 ``tracer_provider``: a TracerProvider, defaults to global
-                ``request_hook``: A hook that receives the span and request that is called
-                    right after the span is created
-                ``response_hook``: A hook that receives the span, request, and response
-                    that is called right before the span ends
+                ``request_hook``: A ``httpx.Client`` hook that receives the span and request
+                    that is called right after the span is created
+                ``response_hook``: A ``httpx.Client`` hook that receives the span, request,
+                    and response that is called right before the span ends
+                ``async_request_hook``: Async ``request_hook`` for ``httpx.AsyncClient``
+                ``async_response_hook``: Async``response_hook`` for ``httpx.AsyncClient``
         """
         self._original_client = httpx.Client
         self._original_async_client = httpx.AsyncClient
         request_hook = kwargs.get("request_hook")
         response_hook = kwargs.get("response_hook")
+        async_request_hook = kwargs.get("async_request_hook", request_hook)
+        async_response_hook = kwargs.get("async_response_hook", response_hook)
         if callable(request_hook):
             _InstrumentedClient._request_hook = request_hook
-            _InstrumentedAsyncClient._request_hook = request_hook
+        if callable(async_request_hook):
+            _InstrumentedAsyncClient._request_hook = async_request_hook
         if callable(response_hook):
             _InstrumentedClient._response_hook = response_hook
-            _InstrumentedAsyncClient._response_hook = response_hook
+        if callable(async_response_hook):
+            _InstrumentedAsyncClient._response_hook = async_response_hook
         tracer_provider = kwargs.get("tracer_provider")
         _InstrumentedClient._tracer_provider = tracer_provider
         _InstrumentedAsyncClient._tracer_provider = tracer_provider
@@ -546,8 +582,12 @@ def _uninstrument(self, **kwargs):
     def instrument_client(
         client: typing.Union[httpx.Client, httpx.AsyncClient],
         tracer_provider: TracerProvider = None,
-        request_hook: typing.Optional[RequestHook] = None,
-        response_hook: typing.Optional[ResponseHook] = None,
+        request_hook: typing.Union[
+            typing.Optional[RequestHook], typing.Optional[AsyncRequestHook]
+        ] = None,
+        response_hook: typing.Union[
+            typing.Optional[ResponseHook], typing.Optional[AsyncResponseHook]
+        ] = None,
     ) -> None:
         """Instrument httpx Client or AsyncClient
 

diff --git a/instrumentation/opentelemetry-instrumentation-httpx/tests/test_httpx_integration.py b/instrumentation/opentelemetry-instrumentation-httpx/tests/test_httpx_integration.py
@@ -421,6 +421,28 @@ def test_response_hook(self):
             )
             HTTPXClientInstrumentor().uninstrument()
 
+        def test_response_hook_sync_async_kwargs(self):
+            HTTPXClientInstrumentor().instrument(
+                tracer_provider=self.tracer_provider,
+                response_hook=_response_hook,
+                async_response_hook=_async_response_hook,
+            )
+            client = self.create_client()
+            result = self.perform_request(self.URL, client=client)
+
+            self.assertEqual(result.text, "Hello!")
+            span = self.assert_span()
+            self.assertEqual(
+                span.attributes,
+                {
+                    SpanAttributes.HTTP_METHOD: "GET",
+                    SpanAttributes.HTTP_URL: self.URL,
+                    SpanAttributes.HTTP_STATUS_CODE: 200,
+                    HTTP_RESPONSE_BODY: "Hello!",
+                },
+            )
+            HTTPXClientInstrumentor().uninstrument()
+
         def test_request_hook(self):
             HTTPXClientInstrumentor().instrument(
                 tracer_provider=self.tracer_provider,
@@ -434,6 +456,20 @@ def test_request_hook(self):
             self.assertEqual(span.name, "GET" + self.URL)
             HTTPXClientInstrumentor().uninstrument()
 
+        def test_request_hook_sync_async_kwargs(self):
+            HTTPXClientInstrumentor().instrument(
+                tracer_provider=self.tracer_provider,
+                request_hook=_request_hook,
+                async_request_hook=_async_request_hook,
+            )
+            client = self.create_client()
+            result = self.perform_request(self.URL, client=client)
+
+            self.assertEqual(result.text, "Hello!")
+            span = self.assert_span()
+            self.assertEqual(span.name, "GET" + self.URL)
+            HTTPXClientInstrumentor().uninstrument()
+
         def test_request_hook_no_span_update(self):
             HTTPXClientInstrumentor().instrument(
                 tracer_provider=self.tracer_provider,