From 8efdccdeee079ea37710af8ecb62c74a2639153f Mon Sep 17 00:00:00 2001 From: Peter Giacomo Lombardo Date: Wed, 8 Jan 2020 16:12:39 +0100 Subject: [PATCH] urllib3: Capture response headers when requested --- instana/instrumentation/urllib3.py | 27 ++++++++++---- tests/apps/flaskalino.py | 8 +++- tests/test_urllib3.py | 60 +++++++++++++++++++++++++++++- 3 files changed, 86 insertions(+), 9 deletions(-) diff --git a/instana/instrumentation/urllib3.py b/instana/instrumentation/urllib3.py index 2b8b441a..2d5d992c 100644 --- a/instana/instrumentation/urllib3.py +++ b/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) diff --git a/tests/apps/flaskalino.py b/tests/apps/flaskalino.py index 44137fcf..f883cbe7 100644 --- a/tests/apps/flaskalino.py +++ b/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") diff --git a/tests/test_urllib3.py b/tests/test_urllib3.py index 77b87b5c..accf7ed1 100644 --- a/tests/test_urllib3.py +++ b/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 +