fastapi: capture responseHeadersOnEntrySpans

Signed-off-by: Varsha GS <[email protected]>

instana/instrumentation/asgi.py

@@ -21,11 +21,12 @@ def __init__(self, app):
 
     def _extract_custom_headers(self, span, headers):
         try:
-            for custom_header in agent.options.extra_http_headers:
-                # Headers are in the following format: b'x-header-1'
-                for header_pair in headers:
-                    if header_pair[0].decode('utf-8').lower() == custom_header.lower():
-                        span.set_tag("http.header.%s" % custom_header, header_pair[1].decode('utf-8'))
+            if agent.options.extra_http_headers is not None:
+                for custom_header in agent.options.extra_http_headers:
+                    # Headers are in the following format: b'x-header-1'
+                    for header_pair in headers:
+                        if header_pair[0].decode('utf-8').lower() == custom_header.lower():
+                            span.set_tag("http.header.%s" % custom_header, header_pair[1].decode('utf-8'))
         except Exception:
             logger.debug("extract_custom_headers: ", exc_info=True)
 
@@ -84,6 +85,7 @@ async def send_wrapper(response):
 
                         headers = response.get('headers')
                         if headers is not None:
+                            self._extract_custom_headers(span, headers)
                             async_tracer.inject(span.context, opentracing.Format.BINARY, headers)
                     except Exception:
                         logger.debug("send_wrapper: ", exc_info=True)

tests/apps/fastapi_app/__init__.py

@@ -13,6 +13,6 @@ def launch_fastapi():
     from instana.singletons import agent
 
     # Hack together a manual custom headers list; We'll use this in tests
-    agent.options.extra_http_headers = [u'X-Capture-This', u'X-Capture-That']
+    agent.options.extra_http_headers = [u'X-Capture-This', u'X-Capture-That', u'X-Capture-This-Too']
 
     uvicorn.run(fastapi_server, host='127.0.0.1', port=testenv['fastapi_port'], log_level="critical")

tests/apps/fastapi_app/app.py

@@ -1,7 +1,7 @@
 # (c) Copyright IBM Corp. 2021
 # (c) Copyright Instana Inc. 2020
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Response
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import PlainTextResponse
 from starlette.exceptions import HTTPException as StarletteHTTPException
@@ -24,6 +24,11 @@ async def root():
 async def user(user_id):
     return {"user": user_id}
 
+@fastapi_server.get("/response_headers")
+async def response_headers():
+    headers = {'X-Capture-This-Too': 'this too'}
+    return Response(content=None, headers=headers)
+
 @fastapi_server.get("/400")
 async def four_zero_zero():
     raise HTTPException(status_code=400, detail="400 response")

tests/frameworks/test_fastapi.py

@@ -311,15 +311,15 @@ def test_synthetic_request(server):
     assert (test_span.sy is None)
 
 
-def test_custom_header_capture(server):
-    from instana.singletons import agent
+def test_request_header_capture(server):
 
     # The background FastAPI server is pre-configured with custom headers to capture
 
     request_headers = {
         'X-Capture-This': 'this',
         'X-Capture-That': 'that'
     }
+
     with tracer.start_active_span('test'):
         result = requests.get(testenv["fastapi_server"] + '/', headers=request_headers)
 
@@ -366,3 +366,53 @@ def test_custom_header_capture(server):
     assert ("this" == asgi_span.data["http"]["header"]["X-Capture-This"])
     assert ("X-Capture-That" in asgi_span.data["http"]["header"])
     assert ("that" == asgi_span.data["http"]["header"]["X-Capture-That"])
+
+
+def test_response_header_capture(server):
+
+     # The background FastAPI server is pre-configured with custom headers to capture
+
+    with tracer.start_active_span('test'):
+        result = requests.get(testenv["fastapi_server"] + '/response_headers')
+
+    assert result.status_code == 200
+
+    spans = tracer.recorder.queued_spans()
+    assert len(spans) == 3
+
+    span_filter = lambda span: span.n == "sdk" and span.data['sdk']['name'] == 'test'
+    test_span = get_first_span_by_filter(spans, span_filter)
+    assert (test_span)
+
+    span_filter = lambda span: span.n == "urllib3"
+    urllib3_span = get_first_span_by_filter(spans, span_filter)
+    assert (urllib3_span)
+
+    span_filter = lambda span: span.n == 'asgi'
+    asgi_span = get_first_span_by_filter(spans, span_filter)
+    assert (asgi_span)
+
+    assert (test_span.t == urllib3_span.t == asgi_span.t)
+    assert (asgi_span.p == urllib3_span.s)
+    assert (urllib3_span.p == test_span.s)
+
+    assert "X-INSTANA-T" in result.headers
+    assert result.headers["X-INSTANA-T"] == asgi_span.t
+    assert "X-INSTANA-S" in result.headers
+    assert result.headers["X-INSTANA-S"] == asgi_span.s
+    assert "X-INSTANA-L" in result.headers
+    assert result.headers["X-INSTANA-L"] == '1'
+    assert "Server-Timing" in result.headers
+    assert result.headers["Server-Timing"] == ("intid;desc=%s" % asgi_span.t)
+
+    assert (asgi_span.ec == None)
+    assert (asgi_span.data['http']['host'] == '127.0.0.1')
+    assert (asgi_span.data['http']['path'] == '/response_headers')
+    assert (asgi_span.data['http']['path_tpl'] == '/response_headers')
+    assert (asgi_span.data['http']['method'] == 'GET')
+    assert (asgi_span.data['http']['status'] == 200)
+    assert (asgi_span.data['http']['error'] is None)
+    assert (asgi_span.data['http']['params'] is None)
+
+    assert ("X-Capture-This-Too" in asgi_span.data["http"]["header"])
+    assert ("this too" == asgi_span.data["http"]["header"]["X-Capture-This-Too"])