Skip to content

Commit 5d068a2

Browse files
committed
urllib3: capture requestHeadersOnExitSpans
Signed-off-by: Varsha GS <[email protected]>
1 parent 0170bd4 commit 5d068a2

File tree

2 files changed

+87
-22
lines changed

2 files changed

+87
-22
lines changed

instana/instrumentation/urllib3.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,18 @@
1616
import urllib3
1717

1818

19+
def extract_custom_headers(span, headers):
20+
if agent.options.extra_http_headers is None:
21+
return
22+
try:
23+
for custom_header in agent.options.extra_http_headers:
24+
if custom_header in headers:
25+
span.set_tag("http.header.%s" % custom_header, headers[custom_header])
26+
27+
except Exception:
28+
logger.debug("extract_custom_headers: ", exc_info=True)
29+
30+
1931
def collect(instance, args, kwargs):
2032
""" Build and return a fully qualified URL for this request """
2133
kvs = dict()
@@ -55,10 +67,7 @@ def collect_response(scope, response):
5567
try:
5668
scope.span.set_tag(ext.HTTP_STATUS_CODE, response.status)
5769

58-
if agent.options.extra_http_headers is not None:
59-
for custom_header in agent.options.extra_http_headers:
60-
if custom_header in response.headers:
61-
scope.span.set_tag("http.header.%s" % custom_header, response.headers[custom_header])
70+
extract_custom_headers(scope.span, response.headers)
6271

6372
if 500 <= response.status:
6473
scope.span.mark_as_errored()
@@ -85,6 +94,7 @@ def urlopen_with_instana(wrapped, instance, args, kwargs):
8594
scope.span.set_tag(ext.HTTP_METHOD, kvs['method'])
8695

8796
if 'headers' in kwargs:
97+
extract_custom_headers(scope.span, kwargs['headers'])
8898
active_tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, kwargs['headers'])
8999

90100
response = wrapped(*args, **kwargs)

tests/clients/test_urllib3.py

Lines changed: 73 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,13 @@
22
# (c) Copyright Instana Inc. 2020
33

44
from __future__ import absolute_import
5+
from multiprocessing.pool import ThreadPool
6+
from time import sleep
7+
import unittest
58

69
import urllib3
7-
import unittest
8-
import sys
910
import requests
1011

11-
from multiprocessing.pool import ThreadPool
12-
from time import sleep
13-
1412
import tests.apps.flask_app
1513
from ..helpers import testenv
1614
from instana.singletons import agent, tracer
@@ -81,7 +79,7 @@ def test_get_request(self):
8179
urllib3_span = spans[1]
8280
test_span = spans[2]
8381

84-
assert(r)
82+
self.assertTrue(r)
8583
self.assertEqual(200, r.status)
8684
self.assertIsNone(tracer.active_span)
8785

@@ -128,7 +126,7 @@ def test_get_request_with_query(self):
128126
urllib3_span = spans[1]
129127
test_span = spans[2]
130128

131-
assert(r)
129+
self.assertTrue(r)
132130
self.assertEqual(200, r.status)
133131
self.assertIsNone(tracer.active_span)
134132

@@ -176,7 +174,7 @@ def test_get_request_with_alt_query(self):
176174
urllib3_span = spans[1]
177175
test_span = spans[2]
178176

179-
assert(r)
177+
self.assertTrue(r)
180178
self.assertEqual(200, r.status)
181179
self.assertIsNone(tracer.active_span)
182180

@@ -224,7 +222,7 @@ def test_put_request(self):
224222
urllib3_span = spans[1]
225223
test_span = spans[2]
226224

227-
assert(r)
225+
self.assertTrue(r)
228226
self.assertEqual(404, r.status)
229227
self.assertIsNone(tracer.active_span)
230228

@@ -273,7 +271,7 @@ def test_301_redirect(self):
273271
urllib3_span1 = spans[3]
274272
test_span = spans[4]
275273

276-
assert(r)
274+
self.assertTrue(r)
277275
self.assertEqual(200, r.status)
278276
self.assertIsNone(tracer.active_span)
279277

@@ -345,7 +343,7 @@ def test_302_redirect(self):
345343
urllib3_span1 = spans[3]
346344
test_span = spans[4]
347345

348-
assert(r)
346+
self.assertTrue(r)
349347
self.assertEqual(200, r.status)
350348
self.assertIsNone(tracer.active_span)
351349

@@ -415,7 +413,7 @@ def test_5xx_request(self):
415413
urllib3_span = spans[1]
416414
test_span = spans[2]
417415

418-
assert(r)
416+
self.assertTrue(r)
419417
self.assertEqual(504, r.status)
420418
self.assertIsNone(tracer.active_span)
421419

@@ -478,7 +476,7 @@ def test_exception_logging(self):
478476

479477
wsgi_span, urllib3_span, test_span = spans
480478

481-
assert(r)
479+
self.assertTrue(r)
482480
self.assertEqual(500, r.status)
483481
self.assertIsNone(tracer.active_span)
484482

@@ -569,7 +567,7 @@ def test_requestspkg_get(self):
569567
urllib3_span = spans[1]
570568
test_span = spans[2]
571569

572-
assert(r)
570+
self.assertTrue(r)
573571
self.assertEqual(200, r.status_code)
574572
self.assertIsNone(tracer.active_span)
575573

@@ -619,7 +617,7 @@ def test_requestspkg_get_with_custom_headers(self):
619617
urllib3_span = spans[1]
620618
test_span = spans[2]
621619

622-
assert(r)
620+
self.assertTrue(r)
623621
self.assertEqual(200, r.status_code)
624622
self.assertIsNone(tracer.active_span)
625623

@@ -703,7 +701,7 @@ def test_requestspkg_put(self):
703701

704702
def test_response_header_capture(self):
705703
original_extra_http_headers = agent.options.extra_http_headers
706-
agent.options.extra_http_headers = ['X-Capture-This']
704+
agent.options.extra_http_headers = ['X-Capture-This', 'X-Capture-That']
707705

708706
with tracer.start_active_span('test'):
709707
r = self.http.request('GET', testenv["wsgi_server"] + '/response_headers')
@@ -715,7 +713,7 @@ def test_response_header_capture(self):
715713
urllib3_span = spans[1]
716714
test_span = spans[2]
717715

718-
assert(r)
716+
self.assertTrue(r)
719717
self.assertEqual(200, r.status)
720718
self.assertIsNone(tracer.active_span)
721719

@@ -751,8 +749,65 @@ def test_response_header_capture(self):
751749
self.assertTrue(type(urllib3_span.stack) is list)
752750
self.assertTrue(len(urllib3_span.stack) > 1)
753751

754-
assert "X-Capture-This" in urllib3_span.data["http"]["header"]
752+
self.assertIn("X-Capture-This", urllib3_span.data["http"]["header"])
755753
self.assertEqual("Ok", urllib3_span.data["http"]["header"]["X-Capture-This"])
754+
self.assertIn("X-Capture-That", urllib3_span.data["http"]["header"])
755+
self.assertEqual("Ok too", urllib3_span.data["http"]["header"]["X-Capture-That"])
756756

757757
agent.options.extra_http_headers = original_extra_http_headers
758758

759+
def test_request_header_capture(self):
760+
original_extra_http_headers = agent.options.extra_http_headers
761+
agent.options.extra_http_headers = ['X-Capture-This-Too']
762+
763+
with tracer.start_active_span('test'):
764+
r = self.http.request('GET', testenv["wsgi_server"] + '/',
765+
headers={'X-Capture-This-Too': 'this too'})
766+
767+
spans = self.recorder.queued_spans()
768+
self.assertEqual(3, len(spans))
769+
770+
wsgi_span = spans[0]
771+
urllib3_span = spans[1]
772+
test_span = spans[2]
773+
774+
self.assertTrue(r)
775+
self.assertEqual(200, r.status)
776+
self.assertIsNone(tracer.active_span)
777+
778+
# Same traceId
779+
self.assertEqual(test_span.t, urllib3_span.t)
780+
self.assertEqual(urllib3_span.t, wsgi_span.t)
781+
782+
# Parent relationships
783+
self.assertEqual(urllib3_span.p, test_span.s)
784+
self.assertEqual(wsgi_span.p, urllib3_span.s)
785+
786+
# Error logging
787+
self.assertIsNone(test_span.ec)
788+
self.assertIsNone(urllib3_span.ec)
789+
self.assertIsNone(wsgi_span.ec)
790+
791+
# wsgi
792+
self.assertEqual("wsgi", wsgi_span.n)
793+
self.assertEqual('127.0.0.1:' + str(testenv["wsgi_port"]), wsgi_span.data["http"]["host"])
794+
self.assertEqual('/', wsgi_span.data["http"]["url"])
795+
self.assertEqual('GET', wsgi_span.data["http"]["method"])
796+
self.assertEqual(200, wsgi_span.data["http"]["status"])
797+
self.assertIsNone(wsgi_span.data["http"]["error"])
798+
self.assertIsNone(wsgi_span.stack)
799+
800+
# urllib3
801+
self.assertEqual("test", test_span.data["sdk"]["name"])
802+
self.assertEqual("urllib3", urllib3_span.n)
803+
self.assertEqual(200, urllib3_span.data["http"]["status"])
804+
self.assertEqual(testenv["wsgi_server"] + "/", urllib3_span.data["http"]["url"])
805+
self.assertEqual("GET", urllib3_span.data["http"]["method"])
806+
self.assertIsNotNone(urllib3_span.stack)
807+
self.assertTrue(type(urllib3_span.stack) is list)
808+
self.assertTrue(len(urllib3_span.stack) > 1)
809+
810+
self.assertIn("X-Capture-This-Too", urllib3_span.data["http"]["header"])
811+
self.assertEqual("this too", urllib3_span.data["http"]["header"]["X-Capture-This-Too"])
812+
813+
agent.options.extra_http_headers = original_extra_http_headers

0 commit comments

Comments
 (0)