diff --git a/datadog_lambda/tag_object.py b/datadog_lambda/tag_object.py index 02dc3ebe..151801f6 100644 --- a/datadog_lambda/tag_object.py +++ b/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): diff --git a/tests/test_tag_object.py b/tests/test_tag_object.py index 67622afe..8e5ac3aa 100644 --- a/tests/test_tag_object.py +++ b/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, )