Store span in ASGI scope (#2725) (#2754)

* Store span in ASGI scope

Motivation
----------

In the Django integration we'd like to be able to access the ASGI
span from the scope of the request to modify it with additional tags
which aren't accessible from the span modifier.

* Add no span test case

* Add pytest mark

(cherry picked from commit c5b5bd2)

Co-authored-by: Kyle Verhoog <[email protected]>
Co-authored-by: Tahir H. Butt <[email protected]>

Author: 3 people

ddtrace/contrib/asgi/__init__.py

@@ -58,5 +58,6 @@ def custom_handle_exception_span(exc, span):
 with require_modules(required_modules) as missing_modules:
     if not missing_modules:
         from .middleware import TraceMiddleware
+        from .middleware import span_from_scope
 
-        __all__ = ["TraceMiddleware"]
+        __all__ = ["TraceMiddleware", "span_from_scope"]

ddtrace/contrib/asgi/middleware.py

@@ -1,4 +1,5 @@
 import sys
+from typing import TYPE_CHECKING
 
 import ddtrace
 from ddtrace import config
@@ -12,6 +13,14 @@
 from .utils import guarantee_single_callable
 
 
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import Mapping
+    from typing import Optional
+
+    from ddtrace import Span
+
+
 log = get_logger(__name__)
 
 config._add(
@@ -58,6 +67,12 @@ def _default_handle_exception_span(exc, span):
     span.set_tag(http.STATUS_CODE, 500)
 
 
+def span_from_scope(scope):
+    # type: (Mapping[str, Any]) -> Optional[Span]
+    """Retrieve the top-level ASGI span from the scope."""
+    return scope.get("datadog_span")
+
+
 class TraceMiddleware:
     """
     ASGI application middleware that traces the requests.
@@ -104,6 +119,8 @@ async def __call__(self, scope, receive, send):
             span_type=SpanTypes.WEB,
         )
 
+        scope["datadog_span"] = span
+
         if self.span_modifier:
             self.span_modifier(span, scope)
 

releasenotes/notes/asgi-span-19e772cb2f09ef4d.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    ASGI: store the ASGI span in the scope. The span is stored under the
+      ``"datadog_span"`` key and can be retrieved using the
+      ``ddtrace.contrib.asgi.span_from_scope`` function.

tests/contrib/asgi/test_asgi.py

@@ -8,6 +8,7 @@
 import pytest
 
 from ddtrace.contrib.asgi import TraceMiddleware
+from ddtrace.contrib.asgi import span_from_scope
 from ddtrace.propagation import http as http_propagation
 from tests.utils import DummyTracer
 from tests.utils import override_http_config
@@ -362,3 +363,55 @@ async def test_bad_headers(scope, tracer, test_spans):
         "type": "http.response.body",
         "body": b"*",
     }
+
+
+@pytest.mark.asyncio
+async def test_get_asgi_span(tracer, test_spans):
+    async def test_app(scope, receive, send):
+        message = await receive()
+        if message.get("type") == "http.request":
+            asgi_span = span_from_scope(scope)
+            assert asgi_span is not None
+            assert asgi_span.name == "asgi.request"
+            await send({"type": "http.response.start", "status": 200, "headers": [[b"Content-Type", b"text/plain"]]})
+            await send({"type": "http.response.body", "body": b""})
+
+    app = TraceMiddleware(test_app, tracer=tracer)
+    async with httpx.AsyncClient(app=app) as client:
+        response = await client.get("http://testserver/")
+        assert response.status_code == 200
+
+    with override_http_config("asgi", dict(trace_query_string=True)):
+        app = TraceMiddleware(test_app, tracer=tracer)
+        async with httpx.AsyncClient(app=app) as client:
+            response = await client.get("http://testserver/")
+            assert response.status_code == 200
+
+    async def test_app(scope, receive, send):
+        message = await receive()
+        if message.get("type") == "http.request":
+            root = tracer.current_root_span()
+            assert root.name == "root"
+            asgi_span = span_from_scope(scope)
+            assert asgi_span is not None
+            assert asgi_span.name == "asgi.request"
+            await send({"type": "http.response.start", "status": 200, "headers": [[b"Content-Type", b"text/plain"]]})
+            await send({"type": "http.response.body", "body": b""})
+
+    app = TraceMiddleware(test_app, tracer=tracer)
+    async with httpx.AsyncClient(app=app) as client:
+        with tracer.trace("root"):
+            response = await client.get("http://testserver/")
+            assert response.status_code == 200
+
+    async def test_app_no_middleware(scope, receive, send):
+        message = await receive()
+        if message.get("type") == "http.request":
+            asgi_span = span_from_scope(scope)
+            assert asgi_span is None
+            await send({"type": "http.response.start", "status": 200, "headers": [[b"Content-Type", b"text/plain"]]})
+            await send({"type": "http.response.body", "body": b""})
+
+    async with httpx.AsyncClient(app=test_app_no_middleware) as client:
+        response = await client.get("http://testserver/")
+        assert response.status_code == 200