Django: capture responseHeadersOnEntrySpans

instana/instrumentation/django/middleware.py

@@ -29,19 +29,27 @@ def __init__(self, get_response=None):
         super(InstanaMiddleware, self).__init__(get_response)
         self.get_response = get_response
 
+    def _extract_custom_headers(self, span, headers, format):
+        try:
+            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
+                    django_header = ('HTTP_' + custom_header.upper()).replace('-', '_') if format else custom_header
+
+                    if django_header in headers:
+                        span.set_tag("http.header.%s" % custom_header, headers[django_header])
+
+        except Exception:
+            logger.debug("extract_custom_headers: ", exc_info=True)
+
     def process_request(self, request):
         try:
             env = request.environ
 
             ctx = tracer.extract(ot.Format.HTTP_HEADERS, env)
             request.iscope = tracer.start_active_span('django', child_of=ctx)
 
-            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
-                    django_header = ('HTTP_' + custom_header.upper()).replace('-', '_')
-                    if django_header in env:
-                        request.iscope.span.set_tag("http.header.%s" % custom_header, env[django_header])
+            self._extract_custom_headers(request.iscope.span, env, format=True)
 
             request.iscope.span.set_tag(ext.HTTP_METHOD, request.method)
             if 'PATH_INFO' in env:
@@ -75,7 +83,9 @@ def process_response(self, request, response):
                         path_tpl = None
                 if path_tpl:
                     request.iscope.span.set_tag("http.path_tpl", path_tpl)
+
                 request.iscope.span.set_tag(ext.HTTP_STATUS_CODE, response.status_code)
+                self._extract_custom_headers(request.iscope.span, response.headers, format=False)
                 tracer.inject(request.iscope.span.context, ot.Format.HTTP_HEADERS, response)
                 response['Server-Timing'] = "intid;desc=%s" % request.iscope.span.context.trace_id
         except Exception:

tests/apps/app_django.py

@@ -125,10 +125,17 @@ def complex(request):
     return HttpResponse('Stan wuz here!')
 
 
+def response_with_headers(request):
+    response = HttpResponse(content_type='')
+    response['X-Capture-This-Too'] = 'this too'
+    return response
+
+
 urlpatterns = [
     re_path(r'^$', index, name='index'),
     re_path(r'^cause_error$', cause_error, name='cause_error'),
     re_path(r'^another$', another),
     re_path(r'^not_found$', not_found, name='not_found'),
-    re_path(r'^complex$', complex, name='complex')
+    re_path(r'^complex$', complex, name='complex'),
+    re_path(r'^response_with_headers$', response_with_headers, name='response_with_headers')
 ]

tests/frameworks/test_django.py

@@ -261,8 +261,9 @@ def test_complex_request(self):
         self.assertEqual(200, django_span.data["http"]["status"])
         self.assertEqual('^complex$', django_span.data["http"]["path_tpl"])
 
-    def test_custom_header_capture(self):
+    def test_request_header_capture(self):
         # Hack together a manual custom headers list
+        original_extra_http_headers = agent.options.extra_http_headers
         agent.options.extra_http_headers = [u'X-Capture-This', u'X-Capture-That']
 
         request_headers = dict()
@@ -306,6 +307,49 @@ def test_custom_header_capture(self):
         assert "X-Capture-That" in django_span.data["http"]["header"]
         self.assertEqual("that", django_span.data["http"]["header"]["X-Capture-That"])
 
+        agent.options.extra_http_headers = original_extra_http_headers
+
+    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 = [u'X-Capture-This-Too']
+
+        with tracer.start_active_span('test'):
+            response = self.http.request('GET', self.live_server_url + '/response_with_headers')
+
+        assert response
+        self.assertEqual(200, response.status)
+
+        spans = self.recorder.queued_spans()
+        self.assertEqual(3, len(spans))
+
+        test_span = spans[2]
+        urllib3_span = spans[1]
+        django_span = spans[0]
+
+        self.assertEqual("test", test_span.data["sdk"]["name"])
+        self.assertEqual("urllib3", urllib3_span.n)
+        self.assertEqual("django", django_span.n)
+
+        self.assertEqual(test_span.t, urllib3_span.t)
+        self.assertEqual(urllib3_span.t, django_span.t)
+
+        self.assertEqual(urllib3_span.p, test_span.s)
+        self.assertEqual(django_span.p, urllib3_span.s)
+
+        self.assertEqual(None, django_span.ec)
+        self.assertIsNone(django_span.stack)
+
+        self.assertEqual('/response_with_headers', django_span.data["http"]["url"])
+        self.assertEqual('GET', django_span.data["http"]["method"])
+        self.assertEqual(200, django_span.data["http"]["status"])
+        self.assertEqual('^response_with_headers$', django_span.data["http"]["path_tpl"])
+
+        assert "X-Capture-This-Too" in django_span.data["http"]["header"]
+        self.assertEqual("this too", django_span.data["http"]["header"]["X-Capture-This-Too"])
+
+        agent.options.extra_http_headers = original_extra_http_headers
+
     def test_with_incoming_context(self):
         request_headers = dict()
         request_headers['X-INSTANA-T'] = '1'