Make OpenTelemetry transport respect Httpx context managers

Author: sam-mosleh

instrumentation/opentelemetry-instrumentation-httpx/src/opentelemetry/instrumentation/httpx/__init__.py

@@ -162,6 +162,7 @@ def response_hook(span, request, response):
 """
 import logging
 import typing
+from types import TracebackType
 
 import httpx
 
@@ -293,6 +294,18 @@ def __init__(
         self._request_hook = request_hook
         self._response_hook = response_hook
 
+    def __enter__(self) -> "SyncOpenTelemetryTransport":
+        self._transport.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: typing.Optional[typing.Type[BaseException]] = None,
+        exc_value: typing.Optional[BaseException] = None,
+        traceback: typing.Optional[TracebackType] = None,
+    ) -> None:
+        self._transport.__exit__(exc_type, exc_value, traceback)
+
     def handle_request(
         self,
         *args,
@@ -343,6 +356,9 @@ def handle_request(
 
         return response
 
+    def close(self) -> None:
+        self._transport.close()
+
 
 class AsyncOpenTelemetryTransport(httpx.AsyncBaseTransport):
     """Async transport class that will trace all requests made with a client.
@@ -372,6 +388,18 @@ def __init__(
         self._request_hook = request_hook
         self._response_hook = response_hook
 
+    async def __aenter__(self: A) -> A:  # Use generics for subclass support.
+        await self._transport.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: typing.Optional[typing.Type[BaseException]] = None,
+        exc_value: typing.Optional[BaseException] = None,
+        traceback: typing.Optional[TracebackType] = None,
+    ) -> None:
+        await self._transport.__aexit__(exc_type, exc_value, traceback)
+
     async def handle_async_request(
         self, *args, **kwargs
     ) -> typing.Union[
@@ -423,6 +451,9 @@ async def handle_async_request(
 
         return response
 
+    async def aclose(self) -> None:
+        await self._transport.aclose()
+
 
 class _InstrumentedClient(httpx.Client):
     _tracer_provider = None