Merge pull request #491 from instana/flask-header-capture

Flask: capture responseHeadersOnEntrySpans

Author: GSVarsha

instana/instrumentation/flask/common.py

@@ -9,7 +9,7 @@
 import opentracing.ext.tags as ext
 
 from ...log import logger
-from ...singletons import tracer
+from ...singletons import tracer, agent
 
 
 @wrapt.patch_function_wrapper('flask', 'templating._render')
@@ -77,3 +77,17 @@ def handle_user_exception_with_instana(wrapped, instance, argv, kwargs):
         logger.debug("handle_user_exception_with_instana:", exc_info=True)
 
     return response
+
+
+def extract_custom_headers(span, headers, format):
+    if agent.options.extra_http_headers is None:
+        return
+    try:
+        for custom_header in agent.options.extra_http_headers:
+            # Headers are available in this format: HTTP_X_CAPTURE_THIS
+            flask_header = ('HTTP_' + custom_header.upper()).replace('-', '_') if format else custom_header
+            if flask_header in headers:
+                span.set_tag("http.header.%s" % custom_header, headers[flask_header])
+
+    except Exception:
+        logger.debug("extract_custom_headers: ", exc_info=True)

instana/instrumentation/flask/vanilla.py

@@ -13,6 +13,7 @@
 from ...log import logger
 from ...singletons import agent, tracer
 from ...util.secrets import strip_secrets_from_query
+from .common import extract_custom_headers
 
 path_tpl_re = re.compile('<.*>')
 
@@ -25,12 +26,7 @@ def before_request_with_instana(*argv, **kwargs):
         flask.g.scope = tracer.start_active_span('wsgi', child_of=ctx)
         span = flask.g.scope.span
 
-        if agent.options.extra_http_headers is not None:
-            for custom_header in agent.options.extra_http_headers:
-                # Headers are available in this format: HTTP_X_CAPTURE_THIS
-                header = ('HTTP_' + custom_header.upper()).replace('-', '_')
-                if header in env:
-                    span.set_tag("http.header.%s" % custom_header, env[header])
+        extract_custom_headers(span, env, format=True)
 
         span.set_tag(ext.HTTP_METHOD, flask.request.method)
         if 'PATH_INFO' in env:
@@ -68,6 +64,8 @@ def after_request_with_instana(response):
                 span.mark_as_errored()
 
             span.set_tag(ext.HTTP_STATUS_CODE, int(response.status_code))
+            extract_custom_headers(span, response.headers, format=False)
+
             tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, response.headers)
             response.headers.add('Server-Timing', "intid;desc=%s" % scope.span.context.trace_id)
     except:

instana/instrumentation/flask/with_blinker.py

@@ -11,6 +11,7 @@
 from ...log import logger
 from ...util.secrets import strip_secrets_from_query
 from ...singletons import agent, tracer
+from .common import extract_custom_headers
 
 import flask
 from flask import request_started, request_finished, got_request_exception
@@ -28,12 +29,7 @@ def request_started_with_instana(sender, **extra):
         flask.g.scope = tracer.start_active_span('wsgi', child_of=ctx)
         span = flask.g.scope.span
 
-        if agent.options.extra_http_headers is not None:
-            for custom_header in agent.options.extra_http_headers:
-                # Headers are available in this format: HTTP_X_CAPTURE_THIS
-                header = ('HTTP_' + custom_header.upper()).replace('-', '_')
-                if header in env:
-                    span.set_tag("http.header.%s" % custom_header, env[header])
+        extract_custom_headers(span, env, format=True)
 
         span.set_tag(ext.HTTP_METHOD, flask.request.method)
         if 'PATH_INFO' in env:
@@ -68,6 +64,8 @@ def request_finished_with_instana(sender, response, **extra):
                 span.mark_as_errored()
 
             span.set_tag(ext.HTTP_STATUS_CODE, int(response.status_code))
+            extract_custom_headers(span, response.headers, format=False)
+
             tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, response.headers)
             response.headers.add('Server-Timing', "intid;desc=%s" % scope.span.context.trace_id)
     except:

tests/apps/flask_app/app.py

@@ -161,10 +161,11 @@ def render_error():
 
 @app.route("/response_headers")
 def response_headers():
-    resp = Response("Foo bar baz")
-    resp.headers['X-Capture-This'] = 'Ok'
-    return resp
-
+    headers = {
+        'X-Capture-This': 'Ok',
+        'X-Capture-That': 'Ok too'
+    }
+    return Response("Stan wuz here with headers!", headers=headers)
 
 @app.route("/boto3/sqs")
 def boto3_sqs():