diff --git a/instana/instrumentation/boto3_inst.py b/instana/instrumentation/boto3_inst.py index cf6511b5..e0b3df5e 100644 --- a/instana/instrumentation/boto3_inst.py +++ b/instana/instrumentation/boto3_inst.py @@ -8,7 +8,7 @@ import inspect from ..log import logger -from ..singletons import tracer +from ..singletons import tracer, agent from ..util.traceutils import get_active_tracer try: @@ -16,6 +16,17 @@ 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: diff --git a/tests/clients/boto3/test_boto3_s3.py b/tests/clients/boto3/test_boto3_s3.py index 4b6d6005..5970048e 100644 --- a/tests/clients/boto3/test_boto3_s3.py +++ b/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