fix(dbm): support byte string queries (#5188) [backport to 1.9] (#5198)

Database monitoring and APM Linking only supports unicode sql queries.
This PR extends support to sql queries with the type: `bytes`.

When `DD_DBM_PROPAGATION_MODE=full` the following code snippet raises an
excpetion:

```
import psycopg2
import ddtrace

ddtrace.patch(psycopg=True)
conn = psycopg2.connect(**POSTGRES_CONFIG)

# raises -> `TypeError: can only concatenate str (not "bytes") to str`
conn.cursor().execute(b"select 'blah'")
```


## Checklist

- [x] Change(s) are motivated and described in the PR description.

Co-authored-by: Munir Abdinur <[email protected]>
Co-authored-by: Zachary Groves <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: ZStriker19 <[email protected]>

Author: 4 people

ddtrace/contrib/psycopg/patch.py

@@ -18,13 +18,14 @@
 from ...internal.utils.formats import asbool
 from ...internal.utils.version import parse_version
 from ...propagation._database_monitoring import _DBM_Propagator
+from ...propagation._database_monitoring import default_sql_injector as _default_sql_injector
 
 
 def _psycopg2_sql_injector(dbm_comment, sql_statement):
     # type: (str, Composable) -> Composable
     if isinstance(sql_statement, Composable):
         return psycopg2.sql.SQL(dbm_comment) + sql_statement
-    return dbm_comment + sql_statement
+    return _default_sql_injector(dbm_comment, sql_statement)
 
 
 config._add(

ddtrace/propagation/_database_monitoring.py

@@ -1,4 +1,5 @@
 from typing import TYPE_CHECKING
+from typing import Union  # noqa
 
 from ddtrace.internal.logger import get_logger
 from ddtrace.vendor.sqlcommenter import generate_sql_comment as _generate_sql_comment
@@ -25,23 +26,24 @@
 log = get_logger(__name__)
 
 
-def _default_sql_injector(dbm_comment, sql_statement):
-    # type: (str, str) -> str
+def default_sql_injector(dbm_comment, sql_statement):
+    # type: (str, Union[str, bytes]) -> Union[str, bytes]
     try:
+        if isinstance(sql_statement, bytes):
+            return dbm_comment.encode("utf-8", errors="strict") + sql_statement
         return dbm_comment + sql_statement
-    except TypeError:
+    except (TypeError, ValueError):
         log.warning(
             "Linking Database Monitoring profiles to spans is not supported for the following query type: %s. "
             "To disable this feature please set the following environment variable: "
             "DD_DBM_PROPAGATION_MODE=disabled",
             type(sql_statement),
-            exc_info=True,
         )
     return sql_statement
 
 
 class _DBM_Propagator(object):
-    def __init__(self, sql_pos, sql_kw, sql_injector=_default_sql_injector):
+    def __init__(self, sql_pos, sql_kw, sql_injector=default_sql_injector):
         self.sql_pos = sql_pos
         self.sql_kw = sql_kw
         self.sql_injector = sql_injector

releasenotes/notes/dbm-support-sql-queries-with-the-type-bytes-30fb4b75ea2f8cf8.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    dbm: Support sql queries with the type ``byte``.

tests/contrib/psycopg/test_psycopg.py

@@ -397,6 +397,9 @@ def test_postgres_dbm_propagation_tag(self):
         # test string queries
         cursor.execute("select 'str blah'")
         cursor.executemany("select %s", (("str_foo",), ("str_bar",)))
+        # test byte string queries
+        cursor.execute(b"select 'byte str blah'")
+        cursor.executemany(b"select %s", ((b"bstr_foo",), (b"bstr_bar",)))
         # test composed queries
         cursor.execute(SQL("select 'composed_blah'"))
         cursor.executemany(SQL("select %s"), (("composed_foo",), ("composed_bar",)))
@@ -418,9 +421,16 @@ def test_postgres_dbm_propagation_comment(self):
         cursor.execute("select 'blah'")
         cursor.executemany("select %s", (("foo",), ("bar",)))
         dbm_comment = "/*dddbs='postgres',dde='staging',ddps='orders-app',ddpv='v7343437-d7ac743'*/ "
-        # test string queries
         cursor.__wrapped__.execute.assert_called_once_with(dbm_comment + "select 'blah'")
         cursor.__wrapped__.executemany.assert_called_once_with(dbm_comment + "select %s", (("foo",), ("bar",)))
+        # test byte string queries
+        cursor.__wrapped__.reset_mock()
+        cursor.execute(b"select 'blah'")
+        cursor.executemany(b"select %s", ((b"foo",), (b"bar",)))
+        cursor.__wrapped__.execute.assert_called_once_with(dbm_comment.encode() + b"select 'blah'")
+        cursor.__wrapped__.executemany.assert_called_once_with(
+            dbm_comment.encode() + b"select %s", ((b"foo",), (b"bar",))
+        )
         # test composed queries
         cursor.__wrapped__.reset_mock()
         cursor.execute(SQL("select 'blah'"))

tests/internal/test_database_monitoring.py

@@ -1,5 +1,8 @@
+import logging
+
 import pytest
 
+from ddtrace.propagation._database_monitoring import default_sql_injector
 from ddtrace.settings import _database_monitoring
 from tests.utils import override_env
 
@@ -114,3 +117,24 @@ def test_dbm_propagation_full_mode():
 
     # ensure that dbm tag is set (required for full mode)
     assert dbspan.get_tag(_database_monitoring.DBM_TRACE_INJECTED_TAG) == "true"
+
+
+def test_default_sql_injector(caplog):
+    # test sql injection with unicode str
+    dbm_comment = "/*dddbs='orders-db'*/ "
+    str_query = "select * from table;"
+    assert default_sql_injector(dbm_comment, str_query) == "/*dddbs='orders-db'*/ select * from table;"
+
+    # test sql injection with uft-8 byte str query
+    dbm_comment = "/*dddbs='orders-db'*/ "
+    str_query = "select * from table;".encode("utf-8")
+    assert default_sql_injector(dbm_comment, str_query) == b"/*dddbs='orders-db'*/ select * from table;"
+
+    # test sql injection with a non supported type
+    with caplog.at_level(logging.INFO):
+        dbm_comment = "/*dddbs='orders-db'*/ "
+        non_string_object = object()
+        result = default_sql_injector(dbm_comment, non_string_object)
+        assert result == non_string_object
+
+    assert "Linking Database Monitoring profiles to spans is not supported for the following query type:" in caplog.text

tests/snapshots/tests.contrib.psycopg.test_psycopg.test_postgres_dbm_propagation_tag.json

@@ -17,7 +17,7 @@
       "db.user": "postgres",
       "language": "python",
       "out.host": "127.0.0.1",
-      "runtime-id": "6759aca0ce8d4ff6a328ff125e50c51d"
+      "runtime-id": "976be4b4fbee409ba62cdc546761d15a"
     },
     "metrics": {
       "_dd.agent_psr": 1.0,
@@ -27,10 +27,10 @@
       "_sampling_priority_v1": 1,
       "db.row_count": 1,
       "network.destination.port": 5432,
-      "process_id": 56478
+      "process_id": 58527
     },
-    "duration": 2652000,
-    "start": 1669148045047826000
+    "duration": 8128000,
+    "start": 1677184993764288000
   }],
 [
   {
@@ -51,7 +51,7 @@
       "db.user": "postgres",
       "language": "python",
       "out.host": "127.0.0.1",
-      "runtime-id": "6759aca0ce8d4ff6a328ff125e50c51d",
+      "runtime-id": "976be4b4fbee409ba62cdc546761d15a",
       "sql.executemany": "true"
     },
     "metrics": {
@@ -62,16 +62,16 @@
       "_sampling_priority_v1": 1,
       "db.row_count": 2,
       "network.destination.port": 5432,
-      "process_id": 56478
+      "process_id": 58527
     },
-    "duration": 2410000,
-    "start": 1669148045078020000
+    "duration": 5081000,
+    "start": 1677184993788584000
   }],
 [
   {
     "name": "postgres.query",
     "service": "postgres",
-    "resource": "select 'composed_blah'",
+    "resource": "select 'byte str blah'",
     "trace_id": 2,
     "span_id": 1,
     "parent_id": 0,
@@ -86,7 +86,7 @@
       "db.user": "postgres",
       "language": "python",
       "out.host": "127.0.0.1",
-      "runtime-id": "6759aca0ce8d4ff6a328ff125e50c51d"
+      "runtime-id": "976be4b4fbee409ba62cdc546761d15a"
     },
     "metrics": {
       "_dd.agent_psr": 1.0,
@@ -96,10 +96,10 @@
       "_sampling_priority_v1": 1,
       "db.row_count": 1,
       "network.destination.port": 5432,
-      "process_id": 56478
+      "process_id": 58527
     },
-    "duration": 1301000,
-    "start": 1669148045080500000
+    "duration": 2678000,
+    "start": 1677184993793709000
   }],
 [
   {
@@ -120,7 +120,76 @@
       "db.user": "postgres",
       "language": "python",
       "out.host": "127.0.0.1",
-      "runtime-id": "6759aca0ce8d4ff6a328ff125e50c51d",
+      "runtime-id": "976be4b4fbee409ba62cdc546761d15a",
+      "sql.executemany": "true"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "db.row_count": 2,
+      "network.destination.port": 5432,
+      "process_id": 58527
+    },
+    "duration": 4226000,
+    "start": 1677184993796426000
+  }],
+[
+  {
+    "name": "postgres.query",
+    "service": "postgres",
+    "resource": "select 'composed_blah'",
+    "trace_id": 4,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "sql",
+    "error": 0,
+    "meta": {
+      "_dd.dbm_trace_injected": "true",
+      "_dd.p.dm": "-0",
+      "component": "psycopg",
+      "db.application": "None",
+      "db.name": "postgres",
+      "db.user": "postgres",
+      "language": "python",
+      "out.host": "127.0.0.1",
+      "runtime-id": "976be4b4fbee409ba62cdc546761d15a"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.measured": 1,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "db.row_count": 1,
+      "network.destination.port": 5432,
+      "process_id": 58527
+    },
+    "duration": 1378000,
+    "start": 1677184993800707000
+  }],
+[
+  {
+    "name": "postgres.query",
+    "service": "postgres",
+    "resource": "select %s",
+    "trace_id": 5,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "sql",
+    "error": 0,
+    "meta": {
+      "_dd.dbm_trace_injected": "true",
+      "_dd.p.dm": "-0",
+      "component": "psycopg",
+      "db.application": "None",
+      "db.name": "postgres",
+      "db.user": "postgres",
+      "language": "python",
+      "out.host": "127.0.0.1",
+      "runtime-id": "976be4b4fbee409ba62cdc546761d15a",
       "sql.executemany": "true"
     },
     "metrics": {
@@ -131,8 +200,8 @@
       "_sampling_priority_v1": 1,
       "db.row_count": 2,
       "network.destination.port": 5432,
-      "process_id": 56478
+      "process_id": 58527
     },
-    "duration": 1988000,
-    "start": 1669148045081837000
+    "duration": 2705000,
+    "start": 1677184993802144000
   }]]