Repository: DataDog/dd-trace-py

Issue:
Psycopg patching doesn't properly handle execute_values
The `execute_values` extension in psycopg2 composes and executes the query with b-string, even if you passed the query as a string. Below is the full function from psycopg2.extras

```python
def execute_values(cur, sql, argslist, template=None, page_size=100, fetch=False):
    from psycopg2.sql import Composable
    if isinstance(sql, Composable):
        sql = sql.as_string(cur)

    # we can't just use sql % vals because vals is bytes: if sql is bytes
    # there will be some decoding error because of stupid codec used, and Py3
    # doesn't implement % on bytes.
    if not isinstance(sql, bytes):
        sql = sql.encode(_ext.encodings[cur.connection.encoding])
    pre, post = _split_sql(sql)

    result = [] if fetch else None
    for page in _paginate(argslist, page_size=page_size):
        if template is None:
            template = b'(' + b','.join([b'%s'] * len(page[0])) + b')'
        parts = pre[:]
        for args in page:
            parts.append(cur.mogrify(template, args))
            parts.append(b',')
        parts[-1:] = post
        cur.execute(b''.join(parts))
        if fetch:
            result.extend(cur.fetchall())

    return result
```

The problem is that ddtrace assumes that the "resource" added to a span is a string. The result is that when `span.finish()` is called in the datadog lambda handler and it tries to serialize the span to json, it blows up with "TypeError: Object of type bytes is not JSON serializable".  Upon investigation, I discovered that the ddtrace.internal.encoder.py's JSONEncoder just does a simple json.dumps() on all the spans and the `resource` attribute on the span from the using `execute_values` is bytes, not a string.

I think the solution here is simply to update the Psycopg2TracedCursor class to decode the resource from bytes if it is bytes, like this:

```python
class Psycopg2TracedCursor(dbapi.TracedCursor):
    """TracedCursor for psycopg2"""

    def _trace_method(self, method, name, resource, extra_tags, *args, **kwargs):
        # treat psycopg2.sql.Composable resource objects as strings
        if isinstance(resource, Composable):
            resource = resource.as_string(self.__wrapped__)
        # THIS IS THE NEW PART BELOW (next 2 lines)
        if isinstance(resource, bytes):
            resource = resource.decode('utf-8')
        return super(Psycopg2TracedCursor, self)._trace_method(method, name, resource, extra_tags, *args, **kwargs)
```

### Which version of dd-trace-py are you using?
Lambda layer, v50.
### Which version of pip are you using?
n/a

### How can we reproduce your problem?
Use `execute_values` while inside a tracing context. It should have a 100% failure rate.

### What is the result that you get?
A type error when span.finish() is called and the metrics are furnished to DD.

### What is the result that you expected?
It should work as normal, with the resource decoded as a string.



Discussion:
Hey @jfalkenstein thanks for reporting this bug.

I have been looking into this, and it seems the issue is with the `LogWriter` we use with the Lambda layer since it encodes the span as JSON, where using the tracer with the Datadog agent will use msgpack which properly supports `bytes`.

My viewpoint right now is that we should fix the JSON encoder we use instead of putting a fix just in the psycopg2 integration since we could have the same issue with other integrations as well.


I'll work on putting a test case / fix for this together.

Resolving patch:
diff --git a/ddtrace/internal/encoding.py b/ddtrace/internal/encoding.py
--- a/ddtrace/internal/encoding.py
+++ b/ddtrace/internal/encoding.py
@@ -8,6 +8,9 @@
 from ._encoding import ListStringTable
 from ._encoding import MsgpackEncoderV03
 from ._encoding import MsgpackEncoderV05
+from .compat import PY3
+from .compat import binary_type
+from .compat import ensure_text
 from .logger import get_logger
 
 
@@ -48,17 +51,33 @@ def encode(self, obj):
         raise NotImplementedError()
 
 
-class JSONEncoder(_EncoderBase):
+class JSONEncoder(json.JSONEncoder, _EncoderBase):
     content_type = "application/json"
 
     def encode_traces(self, traces):
-        normalized_traces = [[span.to_dict() for span in trace] for trace in traces]
+        normalized_traces = [[JSONEncoder._normalize_span(span.to_dict()) for span in trace] for trace in traces]
         return self.encode(normalized_traces)
 
     @staticmethod
-    def encode(obj):
-        # type: (Any) -> str
-        return json.dumps(obj)
+    def _normalize_span(span):
+        # Ensure all string attributes are actually strings and not bytes
+        # DEV: We are deferring meta/metrics to reduce any performance issues.
+        #      Meta/metrics may still contain `bytes` and have encoding issues.
+        span["resource"] = JSONEncoder._normalize_str(span["resource"])
+        span["name"] = JSONEncoder._normalize_str(span["name"])
+        span["service"] = JSONEncoder._normalize_str(span["service"])
+        return span
+
+    @staticmethod
+    def _normalize_str(obj):
+        if obj is None:
+            return obj
+
+        if PY3:
+            return ensure_text(obj, errors="backslashreplace")
+        elif isinstance(obj, binary_type):
+            return obj.decode("utf-8", errors="replace")
+        return obj
 
 
 class JSONEncoderV2(JSONEncoder):
@@ -77,6 +96,7 @@ def encode_traces(self, traces):
     def _convert_span(span):
         # type: (Span) -> Dict[str, Any]
         sp = span.to_dict()
+        sp = JSONEncoderV2._normalize_span(sp)
         sp["trace_id"] = JSONEncoderV2._encode_id_to_hex(sp.get("trace_id"))
         sp["parent_id"] = JSONEncoderV2._encode_id_to_hex(sp.get("parent_id"))
         sp["span_id"] = JSONEncoderV2._encode_id_to_hex(sp.get("span_id"))


Test patch:
diff --git a/tests/tracer/test_encoders.py b/tests/tracer/test_encoders.py
--- a/tests/tracer/test_encoders.py
+++ b/tests/tracer/test_encoders.py
@@ -161,19 +161,20 @@ class TestEncoders(TestCase):
 
     def test_encode_traces_json(self):
         # test encoding for JSON format
-        traces = []
-        traces.append(
+        traces = [
             [
                 Span(name="client.testing", tracer=None),
                 Span(name="client.testing", tracer=None),
-            ]
-        )
-        traces.append(
+            ],
             [
                 Span(name="client.testing", tracer=None),
                 Span(name="client.testing", tracer=None),
-            ]
-        )
+            ],
+            [
+                Span(name=b"client.testing", tracer=None),
+                Span(name=b"client.testing", tracer=None),
+            ],
+        ]
 
         encoder = JSONEncoder()
         spans = encoder.encode_traces(traces)
@@ -182,28 +183,30 @@ def test_encode_traces_json(self):
         # test the encoded output that should be a string
         # and the output must be flatten
         assert isinstance(spans, string_type)
-        assert len(items) == 2
+        assert len(items) == 3
         assert len(items[0]) == 2
         assert len(items[1]) == 2
-        for i in range(2):
+        assert len(items[2]) == 2
+        for i in range(3):
             for j in range(2):
                 assert "client.testing" == items[i][j]["name"]
 
     def test_encode_traces_json_v2(self):
         # test encoding for JSON format
-        traces = []
-        traces.append(
+        traces = [
             [
                 Span(name="client.testing", tracer=None, span_id=0xAAAAAA),
                 Span(name="client.testing", tracer=None, span_id=0xAAAAAA),
-            ]
-        )
-        traces.append(
+            ],
             [
                 Span(name="client.testing", tracer=None, span_id=0xAAAAAA),
                 Span(name="client.testing", tracer=None, span_id=0xAAAAAA),
-            ]
-        )
+            ],
+            [
+                Span(name=b"client.testing", tracer=None, span_id=0xAAAAAA),
+                Span(name=b"client.testing", tracer=None, span_id=0xAAAAAA),
+            ],
+        ]
 
         encoder = JSONEncoderV2()
         spans = encoder.encode_traces(traces)
@@ -211,10 +214,11 @@ def test_encode_traces_json_v2(self):
         # test the encoded output that should be a string
         # and the output must be flatten
         assert isinstance(spans, string_type)
-        assert len(items) == 2
+        assert len(items) == 3
         assert len(items[0]) == 2
         assert len(items[1]) == 2
-        for i in range(2):
+        assert len(items[2]) == 2
+        for i in range(3):
             for j in range(2):
                 assert "client.testing" == items[i][j]["name"]
                 assert isinstance(items[i][j]["span_id"], string_type)
@@ -235,6 +239,12 @@ def test_encode_traces_msgpack_v03(self):
                 Span(name="client.testing", tracer=None),
             ]
         )
+        encoder.put(
+            [
+                Span(name=b"client.testing", tracer=None),
+                Span(name=b"client.testing", tracer=None),
+            ]
+        )
 
         spans = encoder.encode()
         items = encoder._decode(spans)
@@ -242,10 +252,11 @@ def test_encode_traces_msgpack_v03(self):
         # test the encoded output that should be a string
         # and the output must be flatten
         assert isinstance(spans, msgpack_type)
-        assert len(items) == 2
+        assert len(items) == 3
         assert len(items[0]) == 2
         assert len(items[1]) == 2
-        for i in range(2):
+        assert len(items[2]) == 2
+        for i in range(3):
             for j in range(2):
                 assert b"client.testing" == items[i][j][b"name"]
 
@@ -628,3 +639,53 @@ def run(self):
 
     unpacked = decode(encoder.encode(), reconstruct=True)
     assert unpacked is not None
+
+
+@pytest.mark.parametrize("encoder_cls", ["JSONEncoder", "JSONEncoderV2"])
+def test_json_encoder_traces_bytes(encoder_cls, run_python_code_in_subprocess):
+    """
+    Regression test for: https://github.com/DataDog/dd-trace-py/issues/3115
+
+    Ensure we properly decode `bytes` objects when encoding with the JSONEncoder
+    """
+    # Run test in a subprocess to test without setting file encoding to utf-8
+    code = """
+import json
+
+from ddtrace.internal.compat import PY3
+from ddtrace.internal.encoding import {0}
+from ddtrace.span import Span
+
+encoder = {0}()
+data = encoder.encode_traces(
+         [
+             [
+                 Span(name=b"\\x80span.a", tracer=None),
+                 Span(name=u"\\x80span.b", tracer=None),
+                 Span(name="\\x80span.b", tracer=None),
+             ]
+         ]
+    )
+traces = json.loads(data)
+if "{0}" == "JSONEncoderV2":
+    traces = traces["traces"]
+
+assert len(traces) == 1
+span_a, span_b, span_c = traces[0]
+
+if PY3:
+    assert "\\\\x80span.a" == span_a["name"]
+    assert u"\\x80span.b" == span_b["name"]
+    assert u"\\x80span.b" == span_c["name"]
+else:
+    assert u"\\ufffdspan.a" == span_a["name"]
+    assert u"\\x80span.b" == span_b["name"]
+    assert u"\\ufffdspan.b" == span_c["name"]
+""".format(
+        encoder_cls
+    )
+
+    out, err, status, pid = run_python_code_in_subprocess(code)
+    assert status == 0, err
+    assert out == b""
+    assert err == b""