fix(fastapi): apply tracing to user middlewares in fastapi

ddtrace/contrib/fastapi/patch.py

@@ -42,6 +42,18 @@ def traced_init(wrapped, instance, args, kwargs):
     wrapped(*args, **kwargs)
 
 
+def trace_add_middleware(wrapped, instance, args, kwargs):
+    # Ensures user_middlewares are added after ddtrace TraceMiddleware
+    if instance.user_middleware and instance.user_middleware[0].cls is TraceMiddleware:
+        # Overrides FastApi.add_middlware(). Note - fastapi.applications.FastAPI is the
+        # child class of starlette.applications.Starlette():
+        # https://github.com/encode/starlette/blob/0.13.2/starlette/applications.py#L115
+        instance.user_middleware.insert(1, Middleware(*args, **kwargs))
+        instance.middleware_stack = instance.build_middleware_stack()
+    else:
+        wrapped(*args, **kwargs)
+
+
 async def traced_serialize_response(wrapped, instance, args, kwargs):
     """Wrapper for fastapi.routing.serialize_response function.
 
@@ -74,6 +86,7 @@ def patch():
     setattr(fastapi, "_datadog_patch", True)
     Pin().onto(fastapi)
     _w("fastapi.applications", "FastAPI.__init__", traced_init)
+    _w("fastapi.applications", "FastAPI.add_middleware", trace_add_middleware)
     _w("fastapi.routing", "serialize_response", traced_serialize_response)
 
     # We need to check that Starlette instrumentation hasn't already patched these

releasenotes/notes/fastapi-fix-middlewares-705975a535daaea8.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    fastapi: Previously, custom fastapi middlewares configured after application startup were not traced. This fix ensures that all fastapi middlewares are captured in the `fastapi.request` span.

tests/contrib/fastapi/test_fastapi.py

@@ -52,6 +52,23 @@ def application(tracer):
     yield application
 
 
+@pytest.fixture
+def snapshot_app_with_middleware():
+    fastapi_patch()
+
+    application = app.get_app()
+
+    @application.middleware("http")
+    async def traced_middlware(request, call_next):
+        with ddtrace.tracer.trace("traced_middlware"):
+            response = await call_next(request)
+            return response
+
+    yield application
+
+    fastapi_unpatch()
+
+
 @pytest.fixture
 def client(tracer):
     with TestClient(app.get_app()) as test_client:
@@ -590,3 +607,11 @@ def test_host_header(client, tracer, test_spans, host):
     assert test_spans.spans
     request_span = test_spans.spans[0]
     assert request_span.get_tag("http.url") == "http://%s/asynctask" % (host,)
+
+
+@snapshot()
+def test_tracing_in_middleware(snapshot_app_with_middleware):
+    """Test if fastapi middlewares are traced"""
+    with TestClient(snapshot_app_with_middleware) as test_client:
+        r = test_client.get("/", headers={"sleep": "False"})
+        assert r.status_code == 200

tests/snapshots/tests.contrib.fastapi.test_fastapi.test_tracing_in_middleware.json

@@ -0,0 +1,53 @@
+[[
+  {
+    "name": "fastapi.request",
+    "service": "fastapi",
+    "resource": "GET /",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "web",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "http.method": "GET",
+      "http.status_code": "200",
+      "http.url": "http://testserver/",
+      "http.useragent": "testclient",
+      "http.version": "1.1",
+      "runtime-id": "7433b22c2562484081ca485a65d19945"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 4144
+    },
+    "duration": 1034000,
+    "start": 1669131973327481000
+  },
+     {
+       "name": "traced_middlware",
+       "service": "fastapi",
+       "resource": "traced_middlware",
+       "trace_id": 0,
+       "span_id": 2,
+       "parent_id": 1,
+       "type": "",
+       "error": 0,
+       "duration": 459000,
+       "start": 1669131973327931000
+     },
+        {
+          "name": "fastapi.serialize_response",
+          "service": "fastapi",
+          "resource": "fastapi.serialize_response",
+          "trace_id": 0,
+          "span_id": 3,
+          "parent_id": 2,
+          "type": "",
+          "error": 0,
+          "duration": 24000,
+          "start": 1669131973328198000
+        }]]