boto3: initial code for capturing request headers

instana/instrumentation/boto3_inst.py

@@ -8,14 +8,25 @@
 import inspect
 
 from ..log import logger
-from ..singletons import tracer
+from ..singletons import tracer, agent
 from ..util.traceutils import get_active_tracer
 
 try:
     import opentracing as ot
     import boto3
     from boto3.s3 import inject
 
+    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 lambda_inject_context(payload, scope):
         """
@@ -40,6 +51,8 @@ def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
         # pylint: disable=protected-access
         active_tracer = get_active_tracer()
 
+        print("\nWrapping BaseClient._make_api_call with instrumentation")
+
         # If we're not tracing, just return
         if active_tracer is None:
             return wrapped(*arg_list, **kwargs)
@@ -76,6 +89,8 @@ def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
                         status = http_dict.get('HTTPStatusCode')
                         if status is not None:
                             scope.span.set_tag('http.status_code', status)
+                        headers = http_dict.get('HTTPHeaders')
+                        extract_custom_headers(scope.span, headers)
 
                 return result
             except Exception as exc:

tests/clients/boto3/test_boto3_s3.py

@@ -9,7 +9,7 @@
 
 from moto import mock_s3
 
-from instana.singletons import tracer
+from instana.singletons import tracer, agent
 from ...helpers import get_first_span_by_filter
 
 pwd = os.path.dirname(os.path.abspath(__file__))
@@ -90,7 +90,7 @@ def test_s3_list_buckets(s3):
 
     result = s3.list_buckets()
     assert len(result['Buckets']) == 0
-    assert result['ResponseMetadata']['HTTPStatusCode'] is 200
+    assert result['ResponseMetadata']['HTTPStatusCode'] == 200
 
     spans = tracer.recorder.queued_spans()
     assert len(spans) == 2
@@ -271,3 +271,63 @@ def test_s3_download_file_obj(s3):
     assert boto_span.data['boto3']['reg'] == 'us-east-1'
     assert boto_span.data['http']['method'] == 'POST'
     assert boto_span.data['http']['url'] == 'https://s3.amazonaws.com:443/download_fileobj'
+
+
+def test_response_header_capture(s3):
+
+    original_extra_http_headers = agent.options.extra_http_headers
+    agent.options.extra_http_headers = ['X-Capture-This-Too', 'X-Capture-That-Too']
+
+    # Access the event system on the S3 client
+    event_system = s3.meta.events
+
+    response_headers = {
+        "X-Capture-This-Too": "this too",
+        "X-Capture-That-Too": "that too",
+    }
+
+    # Create a function that modifies the after-call event args.
+    def modify_after_call_args(http_response, parsed, model, **kwargs):
+        parsed['ResponseMetadata']['HTTPHeaders'].update(response_headers)
+
+    # Register the function to an event
+    event_system.register('after-call', modify_after_call_args)
+
+    with tracer.start_active_span('test'):
+        result = s3.create_bucket(Bucket="aws_bucket_name")
+
+    result = s3.list_buckets()
+    assert len(result['Buckets']) == 1
+    assert result['Buckets'][0]['Name'] == 'aws_bucket_name'
+
+    spans = tracer.recorder.queued_spans()
+    assert len(spans) == 2
+
+    filter = lambda span: span.n == "sdk"
+    test_span = get_first_span_by_filter(spans, filter)
+    assert (test_span)
+
+    filter = lambda span: span.n == "boto3"
+    boto_span = get_first_span_by_filter(spans, filter)
+    assert (boto_span)
+
+    assert (boto_span.t == test_span.t)
+    assert (boto_span.p == test_span.s)
+
+    assert (test_span.ec is None)
+    assert (boto_span.ec is None)
+
+    assert boto_span.data['boto3']['op'] == 'CreateBucket'
+    assert boto_span.data['boto3']['ep'] == 'https://s3.amazonaws.com'
+    assert boto_span.data['boto3']['reg'] == 'us-east-1'
+    assert boto_span.data['boto3']['payload'] == {'Bucket': 'aws_bucket_name'}
+    assert boto_span.data['http']['status'] == 200
+    assert boto_span.data['http']['method'] == 'POST'
+    assert boto_span.data['http']['url'] == 'https://s3.amazonaws.com:443/CreateBucket'
+
+    assert ("X-Capture-This-Too" in boto_span.data["http"]["header"])
+    assert ("this too" == boto_span.data["http"]["header"]["X-Capture-This-Too"])
+    assert ("X-Capture-That-Too" in boto_span.data["http"]["header"])
+    assert ("that too" == boto_span.data["http"]["header"]["X-Capture-That-Too"])
+
+    agent.options.extra_http_headers = original_extra_http_headers