fix(dbm): support byte string queries [backports #5188 to 1.8] (#5445)

Author: mabdinur

ddtrace/contrib/dbapi/__init__.py

@@ -2,6 +2,7 @@
 Generic dbapi tracing code.
 """
 from typing import TYPE_CHECKING
+from typing import Union
 
 import six
 
@@ -181,15 +182,17 @@ def _propagate_dbm_context(self, dbspan, args):
         return args
 
     def _dbm_sql_injector(self, 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
 

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/dbapi/test_dbapi.py

@@ -1,3 +1,5 @@
+import logging
+
 import mock
 import pytest
 
@@ -20,6 +22,10 @@ def setUp(self):
         super(TestTracedCursor, self).setUp()
         self.cursor = mock.Mock()
 
+    @pytest.fixture(autouse=True)
+    def prepare_fixture(self, caplog):
+        self._caplog = caplog
+
     def test_execute_wrapped_is_called_and_returned(self):
         cursor = self.cursor
         cursor.rowcount = 0
@@ -69,6 +75,29 @@ def test_cursor_execute_with_dbm_injection(self):
         # DBM comment should not be added procedure names
         cursor.callproc.assert_called_once_with("procedure_named_moon")
 
+    def test_default_sql_injector(self):
+        # test sql injection with unicode str
+        cursor = TracedCursor(self.cursor, Pin(), {})
+        dbm_comment = "/*dddbs='orders-db'*/ "
+        str_query = "select * from table;"
+        assert cursor._dbm_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 cursor._dbm_sql_injector(dbm_comment, str_query) == b"/*dddbs='orders-db'*/ select * from table;"
+
+        # test sql injection with a non supported type
+        with self._caplog.at_level(logging.INFO):
+            dbm_comment = "/*dddbs='orders-db'*/ "
+            non_string_object = object()
+            result = cursor._dbm_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 self._caplog.text
+            )
+
     def test_executemany_wrapped_is_called_and_returned(self):
         cursor = self.cursor
         cursor.rowcount = 0

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/snapshots/tests.contrib.psycopg.test_psycopg.test_postgres_dbm_propagation_tag.json

@@ -71,7 +71,7 @@
   {
     "name": "postgres.query",
     "service": "postgres",
-    "resource": "select 'composed_blah'",
+    "resource": "select 'byte str blah'",
     "trace_id": 2,
     "span_id": 1,
     "parent_id": 0,
@@ -133,6 +133,75 @@
       "process_id": 56478,
       "sql.rows": 2
     },
-    "duration": 1988000,
-    "start": 1669148045081837000
+    "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",
+      "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.rowcount": 1,
+      "out.port": 5432,
+      "process_id": 56478,
+      "sql.rows": 1
+    },
+    "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",
+      "out.host": "127.0.0.1",
+      "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.rowcount": 2,
+      "out.port": 5432,
+      "process_id": 56478,
+      "sql.rows": 2
+    },
+    "duration": 2705000,
+    "start": 1677184993802144000
   }]]