Pyramid: capture responseHeadersOnEntrySpans

instana/instrumentation/pyramid/tweens.py

@@ -19,6 +19,17 @@ class InstanaTweenFactory(object):
     def __init__(self, handler, registry):
         self.handler = handler
 
+    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:
+                if custom_header in headers:
+                    span.set_tag("http.header.%s" % custom_header, headers[custom_header])
+
+        except Exception:
+            logger.debug("extract_custom_headers: ", exc_info=True)
+
     def __call__(self, request):
         ctx = tracer.extract(ot.Format.HTTP_HEADERS, dict(request.headers))
         scope = tracer.start_active_span('http', child_of=ctx)
@@ -31,13 +42,8 @@ def __call__(self, request):
         if request.matched_route is not None:
             scope.span.set_tag("http.path_tpl", request.matched_route.pattern)
 
-        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
-                h = ('HTTP_' + custom_header.upper()).replace('-', '_')
-                if h in request.headers:
-                    scope.span.set_tag("http.header.%s" % custom_header, request.headers[h])
-
+        self._extract_custom_headers(scope.span, request.headers)
+
         if len(request.query_string):
             scrubbed_params = strip_secrets_from_query(request.query_string, agent.options.secrets_matcher,
                                                        agent.options.secrets_list)
@@ -47,6 +53,8 @@ def __call__(self, request):
         try:
             response = self.handler(request)
 
+            self._extract_custom_headers(scope.span, response.headers)
+
             tracer.inject(scope.span.context, ot.Format.HTTP_HEADERS, response.headers)
             response.headers['Server-Timing'] = "intid;desc=%s" % scope.span.context.trace_id
         except HTTPException as e:

tests/apps/pyramid_app/app.py

@@ -25,6 +25,13 @@ def please_fail(request):
 def tableflip(request):
     raise BaseException("fake exception")
 
+def response_headers(request):
+    headers = {
+        'X-Capture-This': 'Ok',
+        'X-Capture-That': 'Ok too'
+    }
+    return Response("Stan wuz here with headers!", headers=headers)
+
 app = None
 with Configurator() as config:
     config.add_tween('instana.instrumentation.pyramid.tweens.InstanaTweenFactory')
@@ -34,6 +41,8 @@ def tableflip(request):
     config.add_view(please_fail, route_name='fail')
     config.add_route('crash', '/exception')
     config.add_view(tableflip, route_name='crash')
+    config.add_route('response_headers', '/response_headers')
+    config.add_view(response_headers, route_name='response_headers')
     app = config.make_wsgi_app()
 
 pyramid_server = make_server('127.0.0.1', testenv["pyramid_port"], app)

tests/frameworks/test_pyramid.py

@@ -2,13 +2,13 @@
 # (c) Copyright Instana Inc. 2020
 
 from __future__ import absolute_import
-
 import unittest
+
 import urllib3
 
 import tests.apps.pyramid_app
 from ..helpers import testenv
-from instana.singletons import tracer
+from instana.singletons import tracer, agent
 
 
 class TestPyramid(unittest.TestCase):
@@ -40,21 +40,21 @@ def test_get_request(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert response
+        self.assertTrue(response)
         self.assertEqual(200, response.status)
 
-        assert('X-INSTANA-T' in response.headers)
-        assert(int(response.headers['X-INSTANA-T'], 16))
+        self.assertIn('X-INSTANA-T', response.headers)
+        self.assertTrue(int(response.headers['X-INSTANA-T'], 16))
         self.assertEqual(response.headers['X-INSTANA-T'], pyramid_span.t)
 
-        assert('X-INSTANA-S' in response.headers)
-        assert(int(response.headers['X-INSTANA-S'], 16))
+        self.assertIn('X-INSTANA-S', response.headers)
+        self.assertTrue(int(response.headers['X-INSTANA-S'], 16))
         self.assertEqual(response.headers['X-INSTANA-S'], pyramid_span.s)
 
-        assert('X-INSTANA-L' in response.headers)
+        self.assertIn('X-INSTANA-L', response.headers)
         self.assertEqual(response.headers['X-INSTANA-L'], '1')
 
-        assert('Server-Timing' in response.headers)
+        self.assertIn('Server-Timing', response.headers)
         server_timing_value = "intid;desc=%s" % pyramid_span.t
         self.assertEqual(response.headers['Server-Timing'], server_timing_value)
 
@@ -81,17 +81,17 @@ def test_get_request(self):
         # HTTP SDK span
         self.assertEqual("sdk", pyramid_span.n)
 
-        assert(pyramid_span.data["sdk"])
+        self.assertTrue(pyramid_span.data["sdk"])
         self.assertEqual('http', pyramid_span.data["sdk"]["name"])
         self.assertEqual('entry', pyramid_span.data["sdk"]["type"])
 
-        sdk_data = pyramid_span.data["sdk"]["custom"]
-        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_data["tags"]["http.host"])
-        self.assertEqual('/', sdk_data["tags"]["http.url"])
-        self.assertEqual('GET', sdk_data["tags"]["http.method"])
-        self.assertEqual(200, sdk_data["tags"]["http.status"])
-        self.assertNotIn("message", sdk_data["tags"])
-        self.assertNotIn("http.path_tpl", sdk_data["tags"])
+        sdk_custom_tags = pyramid_span.data["sdk"]["custom"]["tags"]
+        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_custom_tags["http.host"])
+        self.assertEqual('/', sdk_custom_tags["http.url"])
+        self.assertEqual('GET', sdk_custom_tags["http.method"])
+        self.assertEqual(200, sdk_custom_tags["http.status"])
+        self.assertNotIn("message", sdk_custom_tags)
+        self.assertNotIn("http.path_tpl", sdk_custom_tags)
 
         # urllib3
         self.assertEqual("test", test_span.data["sdk"]["name"])
@@ -118,7 +118,7 @@ def test_synthetic_request(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert response
+        self.assertTrue(response)
         self.assertEqual(200, response.status)
 
         self.assertTrue(pyramid_span.sy)
@@ -137,21 +137,21 @@ def test_500(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert response
+        self.assertTrue(response)
         self.assertEqual(500, response.status)
 
-        assert('X-INSTANA-T' in response.headers)
-        assert(int(response.headers['X-INSTANA-T'], 16))
+        self.assertIn('X-INSTANA-T', response.headers)
+        self.assertTrue(int(response.headers['X-INSTANA-T'], 16))
         self.assertEqual(response.headers['X-INSTANA-T'], pyramid_span.t)
 
-        assert('X-INSTANA-S' in response.headers)
-        assert(int(response.headers['X-INSTANA-S'], 16))
+        self.assertIn('X-INSTANA-S', response.headers)
+        self.assertTrue(int(response.headers['X-INSTANA-S'], 16))
         self.assertEqual(response.headers['X-INSTANA-S'], pyramid_span.s)
 
-        assert('X-INSTANA-L' in response.headers)
+        self.assertIn('X-INSTANA-L', response.headers)
         self.assertEqual(response.headers['X-INSTANA-L'], '1')
 
-        assert('Server-Timing' in response.headers)
+        self.assertIn('Server-Timing', response.headers)
         server_timing_value = "intid;desc=%s" % pyramid_span.t
         self.assertEqual(response.headers['Server-Timing'], server_timing_value)
 
@@ -175,13 +175,13 @@ def test_500(self):
         self.assertEqual('http', pyramid_span.data["sdk"]["name"])
         self.assertEqual('entry', pyramid_span.data["sdk"]["type"])
 
-        sdk_data = pyramid_span.data["sdk"]["custom"]
-        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_data["tags"]["http.host"])
-        self.assertEqual('/500', sdk_data["tags"]["http.url"])
-        self.assertEqual('GET', sdk_data["tags"]["http.method"])
-        self.assertEqual(500, sdk_data["tags"]["http.status"])
-        self.assertEqual("internal error", sdk_data["tags"]["message"])
-        self.assertNotIn("http.path_tpl", sdk_data["tags"])
+        sdk_custom_tags = pyramid_span.data["sdk"]["custom"]["tags"]
+        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_custom_tags["http.host"])
+        self.assertEqual('/500', sdk_custom_tags["http.url"])
+        self.assertEqual('GET', sdk_custom_tags["http.method"])
+        self.assertEqual(500, sdk_custom_tags["http.status"])
+        self.assertEqual("internal error", sdk_custom_tags["message"])
+        self.assertNotIn("http.path_tpl", sdk_custom_tags)
 
         # urllib3
         self.assertEqual("test", test_span.data["sdk"]["name"])
@@ -205,7 +205,7 @@ def test_exception(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert response
+        self.assertTrue(response)
         self.assertEqual(500, response.status)
 
         self.assertIsNone(tracer.active_span)
@@ -228,13 +228,13 @@ def test_exception(self):
         self.assertEqual('http', pyramid_span.data["sdk"]["name"])
         self.assertEqual('entry', pyramid_span.data["sdk"]["type"])
 
-        sdk_data = pyramid_span.data["sdk"]["custom"]
-        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_data["tags"]["http.host"])
-        self.assertEqual('/exception', sdk_data["tags"]["http.url"])
-        self.assertEqual('GET', sdk_data["tags"]["http.method"])
-        self.assertEqual(500, sdk_data["tags"]["http.status"])
-        self.assertEqual("fake exception", sdk_data["tags"]["message"])
-        self.assertNotIn("http.path_tpl", sdk_data["tags"])
+        sdk_custom_tags = pyramid_span.data["sdk"]["custom"]["tags"]
+        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_custom_tags["http.host"])
+        self.assertEqual('/exception', sdk_custom_tags["http.url"])
+        self.assertEqual('GET', sdk_custom_tags["http.method"])
+        self.assertEqual(500, sdk_custom_tags["http.status"])
+        self.assertEqual("fake exception", sdk_custom_tags["message"])
+        self.assertNotIn("http.path_tpl", sdk_custom_tags)
 
         # urllib3
         self.assertEqual("test", test_span.data["sdk"]["name"])
@@ -245,3 +245,146 @@ def test_exception(self):
         self.assertIsNotNone(urllib3_span.stack)
         self.assertTrue(type(urllib3_span.stack) is list)
         self.assertTrue(len(urllib3_span.stack) > 1)
+
+    def test_response_header_capture(self):
+        # Hack together a manual custom headers list
+        original_extra_http_headers = agent.options.extra_http_headers
+        agent.options.extra_http_headers = ["X-Capture-This", "X-Capture-That"]
+
+        with tracer.start_active_span('test'):
+            response = self.http.request('GET', testenv["pyramid_server"] + '/response_headers')
+
+        spans = self.recorder.queued_spans()
+        self.assertEqual(3, len(spans))
+
+        pyramid_span = spans[0]
+        urllib3_span = spans[1]
+        test_span = spans[2]
+
+        self.assertTrue(response)
+        self.assertEqual(200, response.status)
+
+        # Same traceId
+        self.assertEqual(test_span.t, urllib3_span.t)
+        self.assertEqual(urllib3_span.t, pyramid_span.t)
+
+        # Parent relationships
+        self.assertEqual(urllib3_span.p, test_span.s)
+        self.assertEqual(pyramid_span.p, urllib3_span.s)
+
+        # Synthetic
+        self.assertIsNone(pyramid_span.sy)
+        self.assertIsNone(urllib3_span.sy)
+        self.assertIsNone(test_span.sy)
+
+        # Error logging
+        self.assertIsNone(test_span.ec)
+        self.assertIsNone(urllib3_span.ec)
+        self.assertIsNone(pyramid_span.ec)
+
+        # HTTP SDK span
+        self.assertEqual("sdk", pyramid_span.n)
+
+        self.assertTrue(pyramid_span.data["sdk"])
+        self.assertEqual('http', pyramid_span.data["sdk"]["name"])
+        self.assertEqual('entry', pyramid_span.data["sdk"]["type"])
+
+        sdk_custom_tags = pyramid_span.data["sdk"]["custom"]["tags"]
+        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_custom_tags["http.host"])
+        self.assertEqual('/response_headers', sdk_custom_tags["http.url"])
+        self.assertEqual('GET', sdk_custom_tags["http.method"])
+        self.assertEqual(200, sdk_custom_tags["http.status"])
+        self.assertNotIn("message", sdk_custom_tags)
+
+        # urllib3
+        self.assertEqual("test", test_span.data["sdk"]["name"])
+        self.assertEqual("urllib3", urllib3_span.n)
+        self.assertEqual(200, urllib3_span.data["http"]["status"])
+        self.assertEqual(testenv["pyramid_server"] + '/response_headers', urllib3_span.data["http"]["url"])
+        self.assertEqual("GET", urllib3_span.data["http"]["method"])
+        self.assertIsNotNone(urllib3_span.stack)
+        self.assertTrue(type(urllib3_span.stack) is list)
+        self.assertTrue(len(urllib3_span.stack) > 1)
+
+
+        self.assertTrue(sdk_custom_tags["http.header.X-Capture-This"])
+        self.assertEqual("Ok", sdk_custom_tags["http.header.X-Capture-This"])
+        self.assertTrue(sdk_custom_tags["http.header.X-Capture-That"])
+        self.assertEqual("Ok too", sdk_custom_tags["http.header.X-Capture-That"])
+
+        agent.options.extra_http_headers = original_extra_http_headers
+
+    def test_request_header_capture(self):
+        original_extra_http_headers = agent.options.extra_http_headers
+        agent.options.extra_http_headers = ["X-Capture-This-Too", "X-Capture-That-Too"]
+
+        request_headers = {
+            "X-Capture-This-Too": "this too",
+            "X-Capture-That-Too": "that too",
+        }
+
+        with tracer.start_active_span("test"):
+            response = self.http.request(
+                "GET", testenv["pyramid_server"] + "/", headers=request_headers
+            )
+
+        spans = self.recorder.queued_spans()
+        self.assertEqual(3, len(spans))
+
+        pyramid_span = spans[0]
+        urllib3_span = spans[1]
+        test_span = spans[2]
+
+        self.assertTrue(response)
+        self.assertEqual(200, response.status)
+
+        # Same traceId
+        self.assertEqual(test_span.t, urllib3_span.t)
+        self.assertEqual(urllib3_span.t, pyramid_span.t)
+
+        # Parent relationships
+        self.assertEqual(urllib3_span.p, test_span.s)
+        self.assertEqual(pyramid_span.p, urllib3_span.s)
+
+        # Synthetic
+        self.assertIsNone(pyramid_span.sy)
+        self.assertIsNone(urllib3_span.sy)
+        self.assertIsNone(test_span.sy)
+
+        # Error logging
+        self.assertIsNone(test_span.ec)
+        self.assertIsNone(urllib3_span.ec)
+        self.assertIsNone(pyramid_span.ec)
+
+        # HTTP SDK span
+        self.assertEqual("sdk", pyramid_span.n)
+
+        self.assertTrue(pyramid_span.data["sdk"])
+        self.assertEqual('http', pyramid_span.data["sdk"]["name"])
+        self.assertEqual('entry', pyramid_span.data["sdk"]["type"])
+
+        sdk_custom_tags = pyramid_span.data["sdk"]["custom"]["tags"]
+        self.assertEqual('127.0.0.1:' + str(testenv['pyramid_port']), sdk_custom_tags["http.host"])
+        self.assertEqual('/', sdk_custom_tags["http.url"])
+        self.assertEqual('GET', sdk_custom_tags["http.method"])
+        self.assertEqual(200, sdk_custom_tags["http.status"])
+        self.assertNotIn("message", sdk_custom_tags)
+        self.assertNotIn("http.path_tpl", sdk_custom_tags)
+
+        # urllib3
+        self.assertEqual("test", test_span.data["sdk"]["name"])
+        self.assertEqual("urllib3", urllib3_span.n)
+        self.assertEqual(200, urllib3_span.data["http"]["status"])
+        self.assertEqual(testenv["pyramid_server"] + '/', urllib3_span.data["http"]["url"])
+        self.assertEqual("GET", urllib3_span.data["http"]["method"])
+        self.assertIsNotNone(urllib3_span.stack)
+        self.assertTrue(type(urllib3_span.stack) is list)
+        self.assertTrue(len(urllib3_span.stack) > 1)
+
+        # custom headers
+        self.assertTrue(sdk_custom_tags["http.header.X-Capture-This-Too"])
+        self.assertEqual("this too", sdk_custom_tags["http.header.X-Capture-This-Too"])
+        self.assertTrue(sdk_custom_tags["http.header.X-Capture-That-Too"])
+        self.assertEqual("that too", sdk_custom_tags["http.header.X-Capture-That-Too"])
+
+        agent.options.extra_http_headers = original_extra_http_headers