enhancing tag_object (#367)

* enhancing tag_object

* format

Author: joeyzhao2018

datadog_lambda/tag_object.py

@@ -28,18 +28,27 @@ def tag_object(span, key, obj, depth=0):
             redacted = _redact_val(key, obj[0:5000])
             return span.set_tag(key, redacted)
     if isinstance(obj, int) or isinstance(obj, float) or isinstance(obj, Decimal):
-        return span.set_tag(key, obj)
+        return span.set_tag(key, str(obj))
     if isinstance(obj, list):
         for k, v in enumerate(obj):
             formatted_key = "{}.{}".format(key, k)
             tag_object(span, formatted_key, v, depth)
         return
-    if isinstance(obj, object):
-        for k in obj:
-            v = obj.get(k)
+    if hasattr(obj, "items"):
+        for k, v in obj.items():
             formatted_key = "{}.{}".format(key, k)
             tag_object(span, formatted_key, v, depth)
         return
+    if hasattr(obj, "to_dict"):
+        for k, v in obj.to_dict().items():
+            formatted_key = "{}.{}".format(key, k)
+            tag_object(span, formatted_key, v, depth)
+        return
+    try:
+        value_as_str = str(obj)
+    except Exception:
+        value_as_str = "UNKNOWN"
+    return span.set_tag(key, value_as_str)
 
 
 def _should_try_string(obj):

tests/test_tag_object.py

@@ -19,12 +19,12 @@ def test_tag_object(self):
         tag_object(spanMock, "function.request", payload)
         spanMock.set_tag.assert_has_calls(
             [
-                call("function.request.vals.0.thingOne", 1),
-                call("function.request.vals.1.thingTwo", 2),
+                call("function.request.vals.0.thingOne", "1"),
+                call("function.request.vals.1.thingTwo", "2"),
                 call("function.request.hello", "world"),
                 call("function.request.anotherThing.blah", None),
                 call("function.request.anotherThing.foo", "bar"),
-                call("function.request.anotherThing.nice", True),
+                call("function.request.anotherThing.nice", "True"),
             ],
             True,
         )
@@ -40,12 +40,12 @@ def test_redacted_tag_object(self):
         tag_object(spanMock, "function.request", payload)
         spanMock.set_tag.assert_has_calls(
             [
-                call("function.request.vals.0.thingOne", 1),
-                call("function.request.vals.1.thingTwo", 2),
+                call("function.request.vals.0.thingOne", "1"),
+                call("function.request.vals.1.thingTwo", "2"),
                 call("function.request.authorization", "redacted"),
                 call("function.request.anotherThing.blah", None),
                 call("function.request.anotherThing.password", "redacted"),
-                call("function.request.anotherThing.nice", True),
+                call("function.request.anotherThing.nice", "True"),
             ],
             True,
         )
@@ -62,7 +62,7 @@ def test_json_tag_object(self):
                 call("function.request.token", "redacted"),
                 call("function.request.jsonString.stringifyThisJson.0.here", "is"),
                 call("function.request.jsonString.stringifyThisJson.0.an", "object"),
-                call("function.request.jsonString.stringifyThisJson.0.number", 1),
+                call("function.request.jsonString.stringifyThisJson.0.number", "1"),
             ],
             True,
         )
@@ -79,18 +79,59 @@ def test_unicode_tag_object(self):
                 call("function.request.token", "redacted"),
                 call("function.request.jsonString.stringifyThisJson.0.here", "is"),
                 call("function.request.jsonString.stringifyThisJson.0.an", "object"),
-                call("function.request.jsonString.stringifyThisJson.0.number", 1),
+                call("function.request.jsonString.stringifyThisJson.0.number", "1"),
             ],
             True,
         )
 
     def test_decimal_tag_object(self):
-        payload = {"myValue": Decimal(500.50)}
+        payload = {"myValue": Decimal(500.5)}
         spanMock = MagicMock()
         tag_object(spanMock, "function.request", payload)
         spanMock.set_tag.assert_has_calls(
             [
-                call("function.request.myValue", Decimal(500.50)),
+                call("function.request.myValue", "500.5"),
+            ],
+            True,
+        )
+
+    class CustomResponse(object):
+        """
+        For example, chalice.app.Response class
+        """
+
+        def __init__(self, body, headers=None, status_code: int = 200):
+            self.body = body
+            if headers is None:
+                headers = {}
+            self.headers = headers
+            self.status_code = status_code
+
+        def __str__(self):
+            return str(self.body)
+
+    class ResponseHasToDict(CustomResponse):
+        def to_dict(self):
+            return self.headers
+
+    def test_custom_response(self):
+        payload = self.CustomResponse({"hello": "world"}, {"key1": "val1"}, 200)
+        spanMock = MagicMock()
+        tag_object(spanMock, "function.response", payload)
+        spanMock.set_tag.assert_has_calls(
+            [
+                call("function.response", "{'hello': 'world'}"),
+            ],
+            True,
+        )
+
+    def test_custom_response_to_dict(self):
+        payload = self.ResponseHasToDict({"hello": "world"}, {"key1": "val1"}, 200)
+        spanMock = MagicMock()
+        tag_object(spanMock, "function.response", payload)
+        spanMock.set_tag.assert_has_calls(
+            [
+                call("function.response.key1", "val1"),
             ],
             True,
         )