Merge pull request #488 from instana/fastapi-header-capture

FastAPI: capture responseHeadersOnEntrySpans

Author: GSVarsha

instana/instrumentation/asgi.py

@@ -20,6 +20,8 @@ def __init__(self, app):
         self.app = app
 
     def _extract_custom_headers(self, span, headers):
+        if agent.options.extra_http_headers is None:
+            return        
         try:
             for custom_header in agent.options.extra_http_headers:
                 # Headers are in the following format: b'x-header-1'
@@ -84,6 +86,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

@@ -6,13 +6,24 @@
 from instana.log import logger
 
 testenv["fastapi_port"] = 10816
-testenv["fastapi_server"] = ("http://127.0.0.1:" + str(testenv["fastapi_port"]))
+testenv["fastapi_server"] = "http://127.0.0.1:" + str(testenv["fastapi_port"])
+
 
 def launch_fastapi():
     from .app import fastapi_server
     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 = [
+        "X-Capture-This",
+        "X-Capture-That",
+        "X-Capture-This-Too",
+        "X-Capture-That-Too",
+    ]
 
-    uvicorn.run(fastapi_server, host='127.0.0.1', port=testenv['fastapi_port'], log_level="critical")
+    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,14 @@ 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',
+        'X-Capture-That-Too': 'that too'
+    }
+    return Response("Stan wuz here with headers!", headers=headers)
+
 @fastapi_server.get("/400")
 async def four_zero_zero():
     raise HTTPException(status_code=400, detail="400 response")