Merge pull request #493 from instana/urllib3-header-capture

urllib3: capture requestHeadersOnExitSpans

Author: GSVarsha

instana/instrumentation/urllib3.py

@@ -16,6 +16,18 @@
     import urllib3
 
 
+    def extract_custom_headers(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 collect(instance, args, kwargs):
         """ Build and return a fully qualified URL for this request """
         kvs = dict()
@@ -55,10 +67,7 @@ def collect_response(scope, response):
         try:
             scope.span.set_tag(ext.HTTP_STATUS_CODE, response.status)
 
-            if agent.options.extra_http_headers is not None:
-                for custom_header in agent.options.extra_http_headers:
-                    if custom_header in response.headers:
-                        scope.span.set_tag("http.header.%s" % custom_header, response.headers[custom_header])
+            extract_custom_headers(scope.span, response.headers)
 
             if 500 <= response.status:
                 scope.span.mark_as_errored()
@@ -85,6 +94,7 @@ def urlopen_with_instana(wrapped, instance, args, kwargs):
                     scope.span.set_tag(ext.HTTP_METHOD, kvs['method'])
 
                 if 'headers' in kwargs:
+                    extract_custom_headers(scope.span, kwargs['headers'])
                     active_tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, kwargs['headers'])
 
                 response = wrapped(*args, **kwargs)

tests/clients/test_urllib3.py

@@ -2,15 +2,13 @@
 # (c) Copyright Instana Inc. 2020
 
 from __future__ import absolute_import
+from multiprocessing.pool import ThreadPool
+from time import sleep
+import unittest
 
 import urllib3
-import unittest
-import sys
 import requests
 
-from multiprocessing.pool import ThreadPool
-from time import sleep
-
 import tests.apps.flask_app
 from ..helpers import testenv
 from instana.singletons import agent, tracer
@@ -81,7 +79,7 @@ def test_get_request(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -128,7 +126,7 @@ def test_get_request_with_query(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -176,7 +174,7 @@ def test_get_request_with_alt_query(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -224,7 +222,7 @@ def test_put_request(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(404, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -273,7 +271,7 @@ def test_301_redirect(self):
         urllib3_span1 = spans[3]
         test_span = spans[4]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -345,7 +343,7 @@ def test_302_redirect(self):
         urllib3_span1 = spans[3]
         test_span = spans[4]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -415,7 +413,7 @@ def test_5xx_request(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(504, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -478,7 +476,7 @@ def test_exception_logging(self):
 
         wsgi_span, urllib3_span, test_span = spans
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(500, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -569,7 +567,7 @@ def test_requestspkg_get(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status_code)
         self.assertIsNone(tracer.active_span)
 
@@ -619,7 +617,7 @@ def test_requestspkg_get_with_custom_headers(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status_code)
         self.assertIsNone(tracer.active_span)
 
@@ -703,7 +701,7 @@ def test_requestspkg_put(self):
 
     def test_response_header_capture(self):
         original_extra_http_headers = agent.options.extra_http_headers
-        agent.options.extra_http_headers = ['X-Capture-This']
+        agent.options.extra_http_headers = ['X-Capture-This', 'X-Capture-That']
 
         with tracer.start_active_span('test'):
             r = self.http.request('GET', testenv["wsgi_server"] + '/response_headers')
@@ -715,7 +713,7 @@ def test_response_header_capture(self):
         urllib3_span = spans[1]
         test_span = spans[2]
 
-        assert(r)
+        self.assertTrue(r)
         self.assertEqual(200, r.status)
         self.assertIsNone(tracer.active_span)
 
@@ -751,8 +749,72 @@ def test_response_header_capture(self):
         self.assertTrue(type(urllib3_span.stack) is list)
         self.assertTrue(len(urllib3_span.stack) > 1)
 
-        assert "X-Capture-This" in urllib3_span.data["http"]["header"]
+        self.assertIn("X-Capture-This", urllib3_span.data["http"]["header"])
         self.assertEqual("Ok", urllib3_span.data["http"]["header"]["X-Capture-This"])
+        self.assertIn("X-Capture-That", urllib3_span.data["http"]["header"])
+        self.assertEqual("Ok too", urllib3_span.data["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"):
+            r = self.http.request(
+                "GET", testenv["wsgi_server"] + "/", headers=request_headers
+            )
+
+        spans = self.recorder.queued_spans()
+        self.assertEqual(3, len(spans))
+
+        wsgi_span = spans[0]
+        urllib3_span = spans[1]
+        test_span = spans[2]
+
+        self.assertTrue(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.assertIsNone(test_span.ec)
+        self.assertIsNone(urllib3_span.ec)
+        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('/', 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.assertIsNone(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"] + "/", 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.assertIn("X-Capture-This-Too", urllib3_span.data["http"]["header"])
+        self.assertEqual("this too", urllib3_span.data["http"]["header"]["X-Capture-This-Too"])
+        self.assertIn("X-Capture-That-Too", urllib3_span.data["http"]["header"])
+        self.assertEqual("that too", urllib3_span.data["http"]["header"]["X-Capture-That-Too"])
+
+        agent.options.extra_http_headers = original_extra_http_headers