Add hooks for aiohttp, asgi, starlette, fastAPI, urllib, urllib3 (#576)

Author: Ryo Kather

Diff for: CHANGELOG.md

@@ -6,8 +6,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.4.0-0.23b0...HEAD)
+
+### Added
 - `opentelemetry-sdk-extension-aws` Add AWS resource detectors to extension package
   ([#586](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/586))
+- `opentelemetry-instrumentation-asgi`, `opentelemetry-instrumentation-aiohttp-client`, `openetelemetry-instrumentation-fastapi`,        
+  `opentelemetry-instrumentation-starlette`, `opentelemetry-instrumentation-urllib`, `opentelemetry-instrumentation-urllib3` Added `request_hook` and `response_hook` callbacks 
+  ([#576](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/576))
 
 ## [1.4.0-0.23b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.4.0-0.23b0) - 2021-07-21
 

Diff for: instrumentation/opentelemetry-instrumentation-aiohttp-client/src/opentelemetry/instrumentation/aiohttp_client/__init__.py

@@ -81,13 +81,25 @@ def strip_query_params(url: yarl.URL) -> str:
 )
 from opentelemetry.propagate import inject
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import SpanKind, TracerProvider, get_tracer
+from opentelemetry.trace import Span, SpanKind, TracerProvider, get_tracer
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util.http import remove_url_credentials
 
 _UrlFilterT = typing.Optional[typing.Callable[[str], str]]
-_SpanNameT = typing.Optional[
-    typing.Union[typing.Callable[[aiohttp.TraceRequestStartParams], str], str]
+_RequestHookT = typing.Optional[
+    typing.Callable[[Span, aiohttp.TraceRequestStartParams], None]
+]
+_ResponseHookT = typing.Optional[
+    typing.Callable[
+        [
+            Span,
+            typing.Union[
+                aiohttp.TraceRequestEndParams,
+                aiohttp.TraceRequestExceptionParams,
+            ],
+        ],
+        None,
+    ]
 ]
 
 
@@ -108,7 +120,8 @@ def url_path_span_name(params: aiohttp.TraceRequestStartParams) -> str:
 
 def create_trace_config(
     url_filter: _UrlFilterT = None,
-    span_name: _SpanNameT = None,
+    request_hook: _RequestHookT = None,
+    response_hook: _ResponseHookT = None,
     tracer_provider: TracerProvider = None,
 ) -> aiohttp.TraceConfig:
     """Create an aiohttp-compatible trace configuration.
@@ -134,15 +147,16 @@ def create_trace_config(
         it as a span attribute. This can be useful to remove sensitive data
         such as API keys or user personal information.
 
-    :param str span_name: Override the default span name.
+    :param Callable request_hook: Optional callback that can modify span name and request params.
+    :param Callable response_hook: Optional callback that can modify span name and response params.
     :param tracer_provider: optional TracerProvider from which to get a Tracer
 
     :return: An object suitable for use with :py:class:`aiohttp.ClientSession`.
     :rtype: :py:class:`aiohttp.TraceConfig`
     """
     # `aiohttp.TraceRequestStartParams` resolves to `aiohttp.tracing.TraceRequestStartParams`
-    # which doesn't exist in the aiottp intersphinx inventory.
-    # Explicitly specify the type for the `span_name` param and rtype to work
+    # which doesn't exist in the aiohttp intersphinx inventory.
+    # Explicitly specify the type for the `request_hook` and `response_hook` param and rtype to work
     # around this issue.
 
     tracer = get_tracer(__name__, __version__, tracer_provider)
@@ -161,17 +175,15 @@ async def on_request_start(
             return
 
         http_method = params.method.upper()
-        if trace_config_ctx.span_name is None:
-            request_span_name = "HTTP {}".format(http_method)
-        elif callable(trace_config_ctx.span_name):
-            request_span_name = str(trace_config_ctx.span_name(params))
-        else:
-            request_span_name = str(trace_config_ctx.span_name)
+        request_span_name = "HTTP {}".format(http_method)
 
         trace_config_ctx.span = trace_config_ctx.tracer.start_span(
             request_span_name, kind=SpanKind.CLIENT,
         )
 
+        if callable(request_hook):
+            request_hook(trace_config_ctx.span, params)
+
         if trace_config_ctx.span.is_recording():
             attributes = {
                 SpanAttributes.HTTP_METHOD: http_method,
@@ -198,6 +210,9 @@ async def on_request_end(
         if trace_config_ctx.span is None:
             return
 
+        if callable(response_hook):
+            response_hook(trace_config_ctx.span, params)
+
         if trace_config_ctx.span.is_recording():
             trace_config_ctx.span.set_status(
                 Status(http_status_to_status_code(int(params.response.status)))
@@ -215,6 +230,9 @@ async def on_request_exception(
         if trace_config_ctx.span is None:
             return
 
+        if callable(response_hook):
+            response_hook(trace_config_ctx.span, params)
+
         if trace_config_ctx.span.is_recording() and params.exception:
             trace_config_ctx.span.set_status(Status(StatusCode.ERROR))
             trace_config_ctx.span.record_exception(params.exception)
@@ -223,7 +241,7 @@ async def on_request_exception(
     def _trace_config_ctx_factory(**kwargs):
         kwargs.setdefault("trace_request_ctx", {})
         return types.SimpleNamespace(
-            span_name=span_name, tracer=tracer, url_filter=url_filter, **kwargs
+            tracer=tracer, url_filter=url_filter, **kwargs
         )
 
     trace_config = aiohttp.TraceConfig(
@@ -240,7 +258,8 @@ def _trace_config_ctx_factory(**kwargs):
 def _instrument(
     tracer_provider: TracerProvider = None,
     url_filter: _UrlFilterT = None,
-    span_name: _SpanNameT = None,
+    request_hook: _RequestHookT = None,
+    response_hook: _ResponseHookT = None,
 ):
     """Enables tracing of all ClientSessions
 
@@ -256,7 +275,8 @@ def instrumented_init(wrapped, instance, args, kwargs):
 
         trace_config = create_trace_config(
             url_filter=url_filter,
-            span_name=span_name,
+            request_hook=request_hook,
+            response_hook=response_hook,
             tracer_provider=tracer_provider,
         )
         trace_config._is_instrumented_by_opentelemetry = True
@@ -304,12 +324,14 @@ def _instrument(self, **kwargs):
                 ``url_filter``: A callback to process the requested URL prior to adding
                     it as a span attribute. This can be useful to remove sensitive data
                     such as API keys or user personal information.
-                ``span_name``: Override the default span name.
+                ``request_hook``: An optional callback that is invoked right after a span is created.
+                ``response_hook``: An optional callback which is invoked right before the span is finished processing a response.
         """
         _instrument(
             tracer_provider=kwargs.get("tracer_provider"),
             url_filter=kwargs.get("url_filter"),
-            span_name=kwargs.get("span_name"),
+            request_hook=kwargs.get("request_hook"),
+            response_hook=kwargs.get("response_hook"),
         )
 
     def _uninstrument(self, **kwargs):

Diff for: instrumentation/opentelemetry-instrumentation-aiohttp-client/tests/test_aiohttp_client_integration.py

@@ -33,7 +33,7 @@
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.test_base import TestBase
-from opentelemetry.trace import StatusCode
+from opentelemetry.trace import Span, StatusCode
 
 
 def run_with_test_server(
@@ -161,46 +161,51 @@ def test_not_recording(self):
             self.assertFalse(mock_span.set_attribute.called)
             self.assertFalse(mock_span.set_status.called)
 
-    def test_span_name_option(self):
-        for span_name, method, path, expected in (
-            ("static", "POST", "/static-span-name", "static"),
-            (
-                lambda params: "{} - {}".format(
-                    params.method, params.url.path
-                ),
-                "PATCH",
-                "/some/path",
-                "PATCH - /some/path",
-            ),
+    def test_hooks(self):
+        method = "PATCH"
+        path = "/some/path"
+        expected = "PATCH - /some/path"
+
+        def request_hook(span: Span, params: aiohttp.TraceRequestStartParams):
+            span.update_name("{} - {}".format(params.method, params.url.path))
+
+        def response_hook(
+            span: Span,
+            params: typing.Union[
+                aiohttp.TraceRequestEndParams,
+                aiohttp.TraceRequestExceptionParams,
+            ],
         ):
-            with self.subTest(span_name=span_name, method=method, path=path):
-                host, port = self._http_request(
-                    trace_config=aiohttp_client.create_trace_config(
-                        span_name=span_name
-                    ),
-                    method=method,
-                    url=path,
-                    status_code=HTTPStatus.OK,
-                )
+            span.set_attribute("response_hook_attr", "value")
 
-                self.assert_spans(
-                    [
-                        (
-                            expected,
-                            (StatusCode.UNSET, None),
-                            {
-                                SpanAttributes.HTTP_METHOD: method,
-                                SpanAttributes.HTTP_URL: "http://{}:{}{}".format(
-                                    host, port, path
-                                ),
-                                SpanAttributes.HTTP_STATUS_CODE: int(
-                                    HTTPStatus.OK
-                                ),
-                            },
-                        )
-                    ]
-                )
-                self.memory_exporter.clear()
+        host, port = self._http_request(
+            trace_config=aiohttp_client.create_trace_config(
+                request_hook=request_hook, response_hook=response_hook,
+            ),
+            method=method,
+            url=path,
+            status_code=HTTPStatus.OK,
+        )
+
+        for span in self.memory_exporter.get_finished_spans():
+            self.assertEqual(span.name, expected)
+            self.assertEqual(
+                (span.status.status_code, span.status.description),
+                (StatusCode.UNSET, None),
+            )
+            self.assertEqual(
+                span.attributes[SpanAttributes.HTTP_METHOD], method
+            )
+            self.assertEqual(
+                span.attributes[SpanAttributes.HTTP_URL],
+                "http://{}:{}{}".format(host, port, path),
+            )
+            self.assertEqual(
+                span.attributes[SpanAttributes.HTTP_STATUS_CODE], HTTPStatus.OK
+            )
+            self.assertIn("response_hook_attr", span.attributes)
+            self.assertEqual(span.attributes["response_hook_attr"], "value")
+        self.memory_exporter.clear()
 
     def test_url_filter_option(self):
         # Strips all query params from URL before adding as a span attribute.
@@ -501,19 +506,32 @@ def strip_query_params(url: yarl.URL) -> str:
             span.attributes[SpanAttributes.HTTP_URL],
         )
 
-    def test_span_name(self):
-        def span_name_callback(params: aiohttp.TraceRequestStartParams) -> str:
-            return "{} - {}".format(params.method, params.url.path)
+    def test_hooks(self):
+        def request_hook(span: Span, params: aiohttp.TraceRequestStartParams):
+            span.update_name("{} - {}".format(params.method, params.url.path))
+
+        def response_hook(
+            span: Span,
+            params: typing.Union[
+                aiohttp.TraceRequestEndParams,
+                aiohttp.TraceRequestExceptionParams,
+            ],
+        ):
+            span.set_attribute("response_hook_attr", "value")
 
         AioHttpClientInstrumentor().uninstrument()
-        AioHttpClientInstrumentor().instrument(span_name=span_name_callback)
+        AioHttpClientInstrumentor().instrument(
+            request_hook=request_hook, response_hook=response_hook
+        )
 
         url = "/test-path"
         run_with_test_server(
             self.get_default_request(url), url, self.default_handler
         )
         span = self.assert_spans(1)
         self.assertEqual("GET - /test-path", span.name)
+        self.assertIn("response_hook_attr", span.attributes)
+        self.assertEqual(span.attributes["response_hook_attr"], "value")
 
 
 class TestLoadingAioHttpInstrumentor(unittest.TestCase):