[aws]Flatten span tag names (#768)

* Add blacklist for params argument of function

* Remove unnecessary import

* Add blacklist to aiobotocore

* Check params

* Fix tests

* Revert blacklisting of params

* Standardize tags for botocore, boto, aiobotocore

* Fix flake8

* Remove resource from tag name

* Change bucket name for error

* Add botocore tests to CI

* Pull out function for adding span tags

* Fix typo

* Add flatten_dict

* Use flatten_dict to cleanup code

* Add note on source

* Fix quotes

* Fix deactivated test

* Add comment

* Double quotes

* Fix flake, remove nose

Author: majorgreys

.circleci/config.yml

@@ -134,7 +134,7 @@ jobs:
       - checkout
       - *restore_cache_step
       - run: tox -e 'boto_contrib-{py27,py34}-boto' --result-json /tmp/boto.1.results
-      - run: tox -e 'botocore_contrib-{py27,py34}-botocore' --result-json /tmp/boto.2.results
+      - run: tox -e 'botocore_contrib-{py27,py34,py35,py36}-botocore' --result-json /tmp/boto.2.results
       - persist_to_workspace:
           root: /tmp
           paths:

ddtrace/contrib/aiobotocore/patch.py

@@ -66,16 +66,6 @@ def __aexit__(self, *args, **kwargs):
             return response
 
 
-def truncate_arg_value(value, max_len=1024):
-    """Truncate values which are bytes and greater than `max_len`.
-    Useful for parameters like 'Body' in `put_object` operations.
-    """
-    if isinstance(value, bytes) and len(value) > max_len:
-        return b'...'
-
-    return value
-
-
 @asyncio.coroutine
 def _wrapped_api_call(original_func, instance, args, kwargs):
     pin = Pin.get_from(instance)
@@ -96,12 +86,7 @@ def _wrapped_api_call(original_func, instance, args, kwargs):
             operation = None
             span.resource = endpoint_name
 
-        # add args in TRACED_ARGS if exist to the span
-        if not aws.is_blacklist(endpoint_name):
-            for name, value in aws.unpacking_args(args, ARGS_NAME, TRACED_ARGS):
-                if name == 'params':
-                    value = {k: truncate_arg_value(v) for k, v in value.items()}
-                span.set_tag(name, (value))
+        aws.add_span_arg_tags(span, endpoint_name, args, ARGS_NAME, TRACED_ARGS)
 
         region_name = deep_getattr(instance, 'meta.region_name')
 

ddtrace/contrib/boto/patch.py

@@ -78,12 +78,7 @@ def patched_query_request(original_func, instance, args, kwargs):
         else:
             span.resource = endpoint_name
 
-        # Adding the args in AWS_QUERY_TRACED_ARGS if exist to the span
-        if not aws.is_blacklist(endpoint_name):
-            for arg in aws.unpacking_args(
-                args, AWS_QUERY_ARGS_NAME, AWS_QUERY_TRACED_ARGS
-            ):
-                span.set_tag(arg[0], arg[1])
+        aws.add_span_arg_tags(span, endpoint_name, args, AWS_QUERY_ARGS_NAME, AWS_QUERY_TRACED_ARGS)
 
         # Obtaining region name
         region_name = _get_instance_region_name(instance)
@@ -127,19 +122,14 @@ def patched_auth_request(original_func, instance, args, kwargs):
         span_type=SPAN_TYPE,
     ) as span:
 
-        # Adding the args in AWS_AUTH_TRACED_ARGS if exist to the span
-        if not aws.is_blacklist(endpoint_name):
-            for arg in aws.unpacking_args(
-                args, AWS_AUTH_ARGS_NAME, AWS_AUTH_TRACED_ARGS
-            ):
-                span.set_tag(arg[0], arg[1])
-
         if args:
             http_method = args[0]
             span.resource = "%s.%s" % (endpoint_name, http_method.lower())
         else:
             span.resource = endpoint_name
 
+        aws.add_span_arg_tags(span, endpoint_name, args, AWS_AUTH_ARGS_NAME, AWS_AUTH_TRACED_ARGS)
+
         # Obtaining region name
         region_name = _get_instance_region_name(instance)
 

ddtrace/contrib/botocore/patch.py

@@ -15,9 +15,9 @@
 # Original botocore client class
 _Botocore_client = botocore.client.BaseClient
 
-SPAN_TYPE = "http"
-ARGS_NAME = ("action", "params", "path", "verb")
-TRACED_ARGS = ["params", "path", "verb"]
+SPAN_TYPE = 'http'
+ARGS_NAME = ('action', 'params', 'path', 'verb')
+TRACED_ARGS = ['params', 'path', 'verb']
 
 
 def patch():
@@ -55,10 +55,7 @@ def patched_api_call(original_func, instance, args, kwargs):
         else:
             span.resource = endpoint_name
 
-        # Adding the args in TRACED_ARGS if exist to the span
-        if not aws.is_blacklist(endpoint_name):
-            for arg in aws.unpacking_args(args, ARGS_NAME, TRACED_ARGS):
-                span.set_tag(arg[0], arg[1])
+        aws.add_span_arg_tags(span, endpoint_name, args, ARGS_NAME, TRACED_ARGS)
 
         region_name = deep_getattr(instance, "meta.region_name")
 

ddtrace/ext/aws.py

@@ -1,30 +1,29 @@
+from ..utils.formats import flatten_dict
+
+
 BLACKLIST_ENDPOINT = ["kms", "sts"]
 
 
-def is_blacklist(endpoint_name):
-    """Protecting the args sent to kms, sts to avoid security leaks
-    if kms disabled test_kms_client in test/contrib/botocore  will fail
-    if sts disabled test_sts_client in test/contrib/boto contrib will fail
+def truncate_arg_value(value, max_len=1024):
+    """Truncate values which are bytes and greater than `max_len`.
+    Useful for parameters like 'Body' in `put_object` operations.
     """
-    return endpoint_name in BLACKLIST_ENDPOINT
+    if isinstance(value, bytes) and len(value) > max_len:
+        return b'...'
 
+    return value
 
-def unpacking_args(args, args_name, traced_args_list):
-    """
-    @params:
-        args: tuple of args sent to a patched function
-        args_name: tuple containing the names of all the args that can be sent
-        traced_args_list: list of names of the args we want to trace
-    Returns a list of (arg name, arg) of the args we want to trace
-    The number of args being variable from one call to another, this function
-    will parse t"""
-    index = 0
-    response = []
-    for arg in args:
-        if arg and args_name[index] in traced_args_list:
-            response += [(args_name[index], arg)]
-        index += 1
-    return response
+
+def add_span_arg_tags(span, endpoint_name, args, args_names, args_traced):
+    if endpoint_name not in BLACKLIST_ENDPOINT:
+        tags = dict(
+            (name, value)
+            for (name, value) in zip(args_names, args)
+            if name in args_traced
+        )
+        tags = flatten_dict(tags)
+        tags = {k: truncate_arg_value(v) for k, v in tags.items()}
+        span.set_tags(tags)
 
 
 REGION = "aws.region"

ddtrace/utils/formats.py

@@ -65,3 +65,16 @@ def asbool(value):
         return value
 
     return value.lower() in ("true", "1")
+
+
+def flatten_dict(d, sep='.', prefix=''):
+    """
+    Returns a normalized dict of depth 1 with keys in order of embedding
+
+    """
+    # adapted from https://stackoverflow.com/a/19647596
+    return {
+        prefix + sep + k if prefix else k: v
+        for kk, vv in d.items()
+        for k, v in flatten_dict(vv, sep, kk).items()
+    } if isinstance(d, dict) else {prefix: d}

tests/contrib/aiobotocore/test.py

@@ -4,6 +4,8 @@
 from botocore.errorfactory import ClientError
 
 from ddtrace.contrib.aiobotocore.patch import patch, unpatch
+from ddtrace.ext import http
+from ddtrace.compat import stringify
 
 from .utils import aiobotocore_client
 from ..asyncio.utils import AsyncioTestCase, mark_asyncio
@@ -59,11 +61,33 @@ def test_s3_client(self):
         eq_(span.resource, 's3.listbuckets')
         eq_(span.name, 's3.command')
 
+    @mark_asyncio
+    def test_s3_put(self):
+        params = dict(Key='foo', Bucket='mybucket', Body=b'bar')
+
+        with aiobotocore_client('s3', self.tracer) as s3:
+            yield from s3.create_bucket(Bucket='mybucket')
+            yield from s3.put_object(**params)
+
+        spans = [trace[0] for trace in self.tracer.writer.pop_traces()]
+        assert spans
+        self.assertEqual(len(spans), 2)
+        self.assertEqual(spans[0].get_tag('aws.operation'), 'CreateBucket')
+        self.assertEqual(spans[0].get_tag(http.STATUS_CODE), '200')
+        self.assertEqual(spans[0].service, 'aws.s3')
+        self.assertEqual(spans[0].resource, 's3.createbucket')
+        self.assertEqual(spans[1].get_tag('aws.operation'), 'PutObject')
+        self.assertEqual(spans[1].resource, 's3.putobject')
+        self.assertEqual(spans[1].get_tag('params.Key'), stringify(params['Key']))
+        self.assertEqual(spans[1].get_tag('params.Bucket'), stringify(params['Bucket']))
+        self.assertEqual(spans[1].get_tag('params.Body'), stringify(params['Body']))
+
     @mark_asyncio
     def test_s3_client_error(self):
         with aiobotocore_client('s3', self.tracer) as s3:
             with assert_raises(ClientError):
-                yield from s3.list_objects(Bucket='mybucket')
+                # FIXME: add proper clean-up to tearDown
+                yield from s3.list_objects(Bucket='doesnotexist')
 
         traces = self.tracer.writer.pop_traces()
         eq_(len(traces), 1)