fix failures

DataDog · Jan 22, 2025 · 05b064d · 05b064d
1 parent 2f56812
commit 05b064d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 23 deletions.
diff --git a/ddtrace/contrib/internal/aiohttp/middlewares.py b/ddtrace/contrib/internal/aiohttp/middlewares.py
@@ -35,7 +35,7 @@ async def attach_context(request):
         analytics_enabled = app[CONFIG_KEY]["analytics_enabled"]
         # Create a new context based on the propagated information.
 
-        with core.context_with_data(
+        ctx = core.context_with_data(
             "aiohttp.request",
             span_name=schematize_url_operation("aiohttp.request", protocol="http", direction=SpanDirection.INBOUND),
             span_type=SpanTypes.WEB,
@@ -48,23 +48,25 @@ async def attach_context(request):
             headers_case_sensitive=True,
             analytics_enabled=analytics_enabled,
             analytics_sample_rate=app[CONFIG_KEY].get("analytics_sample_rate", True),
-        ) as ctx, ctx.span as req_span:
-            ctx.set_item("req_span", req_span)
-            core.dispatch("web.request", (ctx, config.aiohttp))
-
-            # attach the context and the root span to the request; the Context
-            # may be freely used by the application code
-            request[REQUEST_CONTEXT_KEY] = req_span.context
-            request[REQUEST_SPAN_KEY] = req_span
-            request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
-            try:
-                response = await handler(request)
-                if isinstance(response, web.StreamResponse):
-                    request.task.add_done_callback(lambda _: finish_request_span(request, response))
-                return response
-            except Exception:
-                req_span.set_traceback()
-                raise
+        )
+        req_span = ctx.span
+
+        ctx.set_item("req_span", req_span)
+        core.dispatch("web.request", (ctx, config.aiohttp))
+
+        # attach the context and the root span to the request; the Context
+        # may be freely used by the application code
+        request[REQUEST_CONTEXT_KEY] = req_span.context
+        request[REQUEST_SPAN_KEY] = req_span
+        request[REQUEST_CONFIG_KEY] = app[CONFIG_KEY]
+        try:
+            response = await handler(request)
+            if isinstance(response, web.StreamResponse):
+                request.task.add_done_callback(lambda _: finish_request_span(request, response))
+            return response
+        except Exception:
+            req_span.set_traceback()
+            raise
 
     return attach_context
 

diff --git a/ddtrace/contrib/internal/cherrypy/middleware.py b/ddtrace/contrib/internal/cherrypy/middleware.py
@@ -75,7 +75,7 @@ def _setup(self):
         cherrypy.request.hooks.attach("after_error_response", self._after_error_response, priority=5)
 
     def _on_start_resource(self):
-        with core.context_with_data(
+        ctx = core.context_with_data(
             "cherrypy.request",
             span_name=SPAN_NAME,
             span_type=SpanTypes.WEB,
@@ -85,11 +85,13 @@ def _on_start_resource(self):
             distributed_headers=cherrypy.request.headers,
             distributed_headers_config=config.cherrypy,
             headers_case_sensitive=True,
-        ) as ctx, ctx.span as req_span:
-            ctx.set_item("req_span", req_span)
-            core.dispatch("web.request", (ctx, config.cherrypy))
+        )
+        req_span = ctx.span
+
+        ctx.set_item("req_span", req_span)
+        core.dispatch("web.request", (ctx, config.cherrypy))
 
-            cherrypy.request._datadog_span = req_span
+        cherrypy.request._datadog_span = req_span
 
     def _after_error_response(self):
         span = getattr(cherrypy.request, "_datadog_span", None)