urllib3: Capture response headers when requested (#210)

Author: pglombardo

instana/instrumentation/urllib3.py

@@ -44,6 +44,23 @@ def collect(instance, args, kwargs):
         else:
             return kvs
 
+    def collect_response(scope, response):
+        try:
+            scope.span.set_tag(ext.HTTP_STATUS_CODE, response.status)
+
+            if agent.extra_headers is not None:
+                for custom_header in agent.extra_headers:
+                    if custom_header in response.headers:
+                        scope.span.set_tag("http.%s" % custom_header, response.headers[custom_header])
+
+            if 500 <= response.status <= 599:
+                scope.span.set_tag("error", True)
+                ec = scope.span.tags.get('ec', 0)
+                scope.span.set_tag("ec", ec + 1)
+        except Exception:
+            logger.debug("collect_response", exc_info=True)
+
+
     @wrapt.patch_function_wrapper('urllib3', 'HTTPConnectionPool.urlopen')
     def urlopen_with_instana(wrapped, instance, args, kwargs):
         parent_span = tracer.active_span
@@ -65,15 +82,11 @@ def urlopen_with_instana(wrapped, instance, args, kwargs):
                 if 'headers' in kwargs:
                     tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, kwargs['headers'])
 
-                rv = wrapped(*args, **kwargs)
+                response = wrapped(*args, **kwargs)
 
-                scope.span.set_tag(ext.HTTP_STATUS_CODE, rv.status)
-                if 500 <= rv.status <= 599:
-                    scope.span.set_tag("error", True)
-                    ec = scope.span.tags.get('ec', 0)
-                    scope.span.set_tag("ec", ec+1)
+                collect_response(scope, response)
 
-                return rv
+                return response
             except Exception as e:
                 scope.span.log_kv({'message': e})
                 scope.span.set_tag("error", True)

tests/apps/flaskalino.py

@@ -3,7 +3,7 @@
 import opentracing.ext.tags as ext
 from flask import Flask, redirect, render_template, render_template_string
 from wsgiref.simple_server import make_server
-from flask import jsonify
+from flask import jsonify, Response
 
 from instana.singletons import tracer
 from ..helpers import testenv
@@ -121,6 +121,12 @@ def render_error():
     return render_template('flask_render_error.html', what='world')
 
 
+@app.route("/response_headers")
+def response_headers():
+    resp = Response("Foo bar baz")
+    resp.headers['X-Capture-This'] = 'Ok'
+    return resp
+
 @app.errorhandler(InvalidUsage)
 def handle_invalid_usage(error):
     logger.error("InvalidUsage error handler invoked")

tests/test_urllib3.py

@@ -5,7 +5,7 @@
 import requests
 import urllib3
 
-from instana.singletons import tracer
+from instana.singletons import agent, tracer
 from .helpers import testenv
 
 
@@ -692,3 +692,61 @@ def test_requestspkg_put(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):
+        original_extra_headers = agent.extra_headers
+        agent.extra_headers = ['X-Capture-This']
+
+        with tracer.start_active_span('test'):
+            r = self.http.request('GET', testenv["wsgi_server"] + '/response_headers')
+
+        spans = self.recorder.queued_spans()
+        self.assertEqual(3, len(spans))
+
+        wsgi_span = spans[0]
+        urllib3_span = spans[1]
+        test_span = spans[2]
+
+        assert(r)
+        self.assertEqual(200, r.status)
+        self.assertIsNone(tracer.active_span)
+
+        # Same traceId
+        self.assertEqual(test_span.t, urllib3_span.t)
+        self.assertEqual(urllib3_span.t, wsgi_span.t)
+
+        # Parent relationships
+        self.assertEqual(urllib3_span.p, test_span.s)
+        self.assertEqual(wsgi_span.p, urllib3_span.s)
+
+        # Error logging
+        self.assertFalse(test_span.error)
+        self.assertIsNone(test_span.ec)
+        self.assertFalse(urllib3_span.error)
+        self.assertIsNone(urllib3_span.ec)
+        self.assertFalse(wsgi_span.error)
+        self.assertIsNone(wsgi_span.ec)
+
+        # wsgi
+        self.assertEqual("wsgi", wsgi_span.n)
+        self.assertEqual('127.0.0.1:' + str(testenv["wsgi_port"]), wsgi_span.data.http.host)
+        self.assertEqual('/response_headers', wsgi_span.data.http.url)
+        self.assertEqual('GET', wsgi_span.data.http.method)
+        self.assertEqual(200, wsgi_span.data.http.status)
+        self.assertIsNone(wsgi_span.data.http.error)
+        self.assertIsNotNone(wsgi_span.stack)
+        self.assertEqual(2, len(wsgi_span.stack))
+
+        # 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["wsgi_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('http.X-Capture-This' in urllib3_span.data.custom.tags)
+
+        agent.extra_headers = original_extra_headers
+