feat(asgi,fastapi,starlette)!: provide both send and receive hooks with scope and message (open-telemetry#2546)

Author: eigenein

CHANGELOG.md

@@ -27,6 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#2425](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2425))
 - `opentelemetry-instrumentation-flask` Add `http.method` to `span.name`
   ([#2454](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2454))
+- ASGI, FastAPI, Starlette: provide both send and receive hooks with `scope` and `message` for internal spans ([#2546](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2546))
 
 ### Added
 

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

@@ -81,15 +81,15 @@ async def hello():
 
 .. code-block:: python
 
-    def server_request_hook(span: Span, scope: dict):
+    def server_request_hook(span: Span, scope: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_request_hook", "some-value")
 
-    def client_request_hook(span: Span, scope: dict):
+    def client_request_hook(span: Span, scope: dict[str, Any], message: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_client_request_hook", "some-value")
 
-    def client_response_hook(span: Span, message: dict):
+    def client_response_hook(span: Span, scope: dict[str, Any], message: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_response_hook", "some-value")
 
@@ -200,6 +200,11 @@ def client_response_hook(span: Span, message: dict):
 from asgiref.compatibility import guarantee_single_callable
 
 from opentelemetry import context, trace
+from opentelemetry.instrumentation.asgi.types import (
+    ClientRequestHook,
+    ClientResponseHook,
+    ServerRequestHook,
+)
 from opentelemetry.instrumentation.asgi.version import __version__  # noqa
 from opentelemetry.instrumentation.propagators import (
     get_global_response_propagator,
@@ -212,7 +217,7 @@ def client_response_hook(span: Span, message: dict):
 from opentelemetry.propagators.textmap import Getter, Setter
 from opentelemetry.semconv.metrics import MetricInstruments
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import Span, set_span_in_context
+from opentelemetry.trace import set_span_in_context
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util.http import (
     OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SANITIZE_FIELDS,
@@ -227,10 +232,6 @@ def client_response_hook(span: Span, message: dict):
     remove_url_credentials,
 )
 
-_ServerRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-_ClientRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-_ClientResponseHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-
 
 class ASGIGetter(Getter[dict]):
     def get(
@@ -454,10 +455,10 @@ class OpenTelemetryMiddleware:
                       Optional: Defaults to get_default_span_details.
         server_request_hook: Optional callback which is called with the server span and ASGI
                       scope object for every incoming request.
-        client_request_hook: Optional callback which is called with the internal span and an ASGI
-                      scope which is sent as a dictionary for when the method receive is called.
-        client_response_hook: Optional callback which is called with the internal span and an ASGI
-                      event which is sent as a dictionary for when the method send is called.
+        client_request_hook: Optional callback which is called with the internal span, and ASGI
+                      scope and event which are sent as dictionaries for when the method receive is called.
+        client_response_hook: Optional callback which is called with the internal span, and ASGI
+                      scope and event which are sent as dictionaries for when the method send is called.
         tracer_provider: The optional tracer provider to use. If omitted
             the current globally configured one is used.
     """
@@ -468,9 +469,9 @@ def __init__(
         app,
         excluded_urls=None,
         default_span_details=None,
-        server_request_hook: _ServerRequestHookT = None,
-        client_request_hook: _ClientRequestHookT = None,
-        client_response_hook: _ClientResponseHookT = None,
+        server_request_hook: ServerRequestHook = None,
+        client_request_hook: ClientRequestHook = None,
+        client_response_hook: ClientResponseHook = None,
         tracer_provider=None,
         meter_provider=None,
         meter=None,
@@ -666,9 +667,9 @@ async def otel_receive():
             with self.tracer.start_as_current_span(
                 " ".join((server_span_name, scope["type"], "receive"))
             ) as receive_span:
-                if callable(self.client_request_hook):
-                    self.client_request_hook(receive_span, scope)
                 message = await receive()
+                if callable(self.client_request_hook):
+                    self.client_request_hook(receive_span, scope, message)
                 if receive_span.is_recording():
                     if message["type"] == "websocket.receive":
                         set_status_code(receive_span, 200)
@@ -691,7 +692,7 @@ async def otel_send(message: dict[str, Any]):
                 " ".join((server_span_name, scope["type"], "send"))
             ) as send_span:
                 if callable(self.client_response_hook):
-                    self.client_response_hook(send_span, message)
+                    self.client_response_hook(send_span, scope, message)
                 if send_span.is_recording():
                     if message["type"] == "http.response.start":
                         status_code = message["status"]

instrumentation/opentelemetry-instrumentation-asgi/src/opentelemetry/instrumentation/asgi/types.py

@@ -0,0 +1,49 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Callable, Dict, Optional
+
+from opentelemetry.trace import Span
+
+_Scope = Dict[str, Any]
+_Message = Dict[str, Any]
+
+ServerRequestHook = Optional[Callable[[Span, _Scope], None]]
+"""
+Incoming request callback type.
+
+Args:
+    - Server span
+    - ASGI scope as a mapping
+"""
+
+ClientRequestHook = Optional[Callable[[Span, _Scope, _Message], None]]
+"""
+Receive callback type.
+
+Args:
+    - Internal span
+    - ASGI scope as a mapping
+    - ASGI event as a mapping
+"""
+
+ClientResponseHook = Optional[Callable[[Span, _Scope, _Message], None]]
+"""
+Send callback type.
+
+Args:
+    - Internal span
+    - ASGI scope as a mapping
+    - ASGI event as a mapping
+"""

instrumentation/opentelemetry-instrumentation-asgi/tests/test_asgi_middleware.py

@@ -683,10 +683,10 @@ def test_hooks(self):
         def server_request_hook(span, scope):
             span.update_name("name from server hook")
 
-        def client_request_hook(recieve_span, request):
-            recieve_span.update_name("name from client request hook")
+        def client_request_hook(receive_span, scope, message):
+            receive_span.update_name("name from client request hook")
 
-        def client_response_hook(send_span, response):
+        def client_response_hook(send_span, scope, message):
             send_span.set_attribute("attr-from-hook", "value")
 
         def update_expected_hook_results(expected):

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

@@ -59,20 +59,20 @@ async def foobar():
 right after a span is created for a request and right before the span is finished for the response.
 
 - The server request hook is passed a server span and ASGI scope object for every incoming request.
-- The client request hook is called with the internal span and an ASGI scope when the method ``receive`` is called.
-- The client response hook is called with the internal span and an ASGI event when the method ``send`` is called.
+- The client request hook is called with the internal span, and ASGI scope and event when the method ``receive`` is called.
+- The client response hook is called with the internal span, and ASGI scope and event when the method ``send`` is called.
 
 .. code-block:: python
 
-    def server_request_hook(span: Span, scope: dict):
+    def server_request_hook(span: Span, scope: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_request_hook", "some-value")
 
-    def client_request_hook(span: Span, scope: dict):
+    def client_request_hook(span: Span, scope: dict[str, Any], message: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_client_request_hook", "some-value")
 
-    def client_response_hook(span: Span, message: dict):
+    def client_response_hook(span: Span, scope: dict[str, Any], message: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_response_hook", "some-value")
 
@@ -172,28 +172,27 @@ def client_response_hook(span: Span, message: dict):
 ---
 """
 import logging
-import typing
 from typing import Collection
 
 import fastapi
 from starlette.routing import Match
 
 from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
+from opentelemetry.instrumentation.asgi.types import (
+    ClientRequestHook,
+    ClientResponseHook,
+    ServerRequestHook,
+)
 from opentelemetry.instrumentation.fastapi.package import _instruments
 from opentelemetry.instrumentation.fastapi.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.metrics import get_meter
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import Span
 from opentelemetry.util.http import get_excluded_urls, parse_excluded_urls
 
 _excluded_urls_from_env = get_excluded_urls("FASTAPI")
 _logger = logging.getLogger(__name__)
 
-_ServerRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-_ClientRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-_ClientResponseHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-
 
 class FastAPIInstrumentor(BaseInstrumentor):
     """An instrumentor for FastAPI
@@ -206,9 +205,9 @@ class FastAPIInstrumentor(BaseInstrumentor):
     @staticmethod
     def instrument_app(
         app: fastapi.FastAPI,
-        server_request_hook: _ServerRequestHookT = None,
-        client_request_hook: _ClientRequestHookT = None,
-        client_response_hook: _ClientResponseHookT = None,
+        server_request_hook: ServerRequestHook = None,
+        client_request_hook: ClientRequestHook = None,
+        client_response_hook: ClientResponseHook = None,
         tracer_provider=None,
         meter_provider=None,
         excluded_urls=None,
@@ -292,9 +291,9 @@ class _InstrumentedFastAPI(fastapi.FastAPI):
     _tracer_provider = None
     _meter_provider = None
     _excluded_urls = None
-    _server_request_hook: _ServerRequestHookT = None
-    _client_request_hook: _ClientRequestHookT = None
-    _client_response_hook: _ClientResponseHookT = None
+    _server_request_hook: ServerRequestHook = None
+    _client_request_hook: ClientRequestHook = None
+    _client_response_hook: ClientResponseHook = None
     _instrumented_fastapi_apps = set()
 
     def __init__(self, *args, **kwargs):

instrumentation/opentelemetry-instrumentation-fastapi/tests/test_fastapi_instrumentation.py

@@ -342,23 +342,23 @@ def server_request_hook(self, span, scope):
         if self._server_request_hook is not None:
             self._server_request_hook(span, scope)
 
-    def client_request_hook(self, receive_span, request):
+    def client_request_hook(self, receive_span, scope, message):
         if self._client_request_hook is not None:
-            self._client_request_hook(receive_span, request)
+            self._client_request_hook(receive_span, scope, message)
 
-    def client_response_hook(self, send_span, response):
+    def client_response_hook(self, send_span, scope, message):
         if self._client_response_hook is not None:
-            self._client_response_hook(send_span, response)
+            self._client_response_hook(send_span, scope, message)
 
     def test_hooks(self):
         def server_request_hook(span, scope):
             span.update_name("name from server hook")
 
-        def client_request_hook(receive_span, request):
+        def client_request_hook(receive_span, scope, message):
             receive_span.update_name("name from client hook")
             receive_span.set_attribute("attr-from-request-hook", "set")
 
-        def client_response_hook(send_span, response):
+        def client_response_hook(send_span, scope, message):
             send_span.update_name("name from response hook")
             send_span.set_attribute("attr-from-response-hook", "value")
 

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

@@ -55,20 +55,22 @@ def home(request):
 right after a span is created for a request and right before the span is finished for the response.
 
 - The server request hook is passed a server span and ASGI scope object for every incoming request.
-- The client request hook is called with the internal span and an ASGI scope when the method ``receive`` is called.
-- The client response hook is called with the internal span and an ASGI event when the method ``send`` is called.
+- The client request hook is called with the internal span, and ASGI scope and event when the method ``receive`` is called.
+- The client response hook is called with the internal span, and ASGI scope and event when the method ``send`` is called.
 
 For example,
 
 .. code-block:: python
 
-    def server_request_hook(span: Span, scope: dict):
+    def server_request_hook(span: Span, scope: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_request_hook", "some-value")
-    def client_request_hook(span: Span, scope: dict):
+
+    def client_request_hook(span: Span, scope: dict[str, Any], message: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_client_request_hook", "some-value")
-    def client_response_hook(span: Span, message: dict):
+
+    def client_response_hook(span: Span, scope: dict[str, Any], message: dict[str, Any]):
         if span and span.is_recording():
             span.set_attribute("custom_user_attribute_from_response_hook", "some-value")
 
@@ -167,27 +169,26 @@ def client_response_hook(span: Span, message: dict):
 API
 ---
 """
-import typing
 from typing import Collection
 
 from starlette import applications
 from starlette.routing import Match
 
 from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
+from opentelemetry.instrumentation.asgi.types import (
+    ClientRequestHook,
+    ClientResponseHook,
+    ServerRequestHook,
+)
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.starlette.package import _instruments
 from opentelemetry.instrumentation.starlette.version import __version__
 from opentelemetry.metrics import get_meter
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import Span
 from opentelemetry.util.http import get_excluded_urls
 
 _excluded_urls = get_excluded_urls("STARLETTE")
 
-_ServerRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-_ClientRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-_ClientResponseHookT = typing.Optional[typing.Callable[[Span, dict], None]]
-
 
 class StarletteInstrumentor(BaseInstrumentor):
     """An instrumentor for starlette
@@ -200,9 +201,9 @@ class StarletteInstrumentor(BaseInstrumentor):
     @staticmethod
     def instrument_app(
         app: applications.Starlette,
-        server_request_hook: _ServerRequestHookT = None,
-        client_request_hook: _ClientRequestHookT = None,
-        client_response_hook: _ClientResponseHookT = None,
+        server_request_hook: ServerRequestHook = None,
+        client_request_hook: ClientRequestHook = None,
+        client_response_hook: ClientResponseHook = None,
         meter_provider=None,
         tracer_provider=None,
     ):
@@ -270,9 +271,9 @@ def _uninstrument(self, **kwargs):
 class _InstrumentedStarlette(applications.Starlette):
     _tracer_provider = None
     _meter_provider = None
-    _server_request_hook: _ServerRequestHookT = None
-    _client_request_hook: _ClientRequestHookT = None
-    _client_response_hook: _ClientResponseHookT = None
+    _server_request_hook: ServerRequestHook = None
+    _client_request_hook: ClientRequestHook = None
+    _client_response_hook: ClientResponseHook = None
     _instrumented_starlette_apps = set()
 
     def __init__(self, *args, **kwargs):