feat: add dbm for all mysql and postgres integrations

ddtrace/contrib/aiomysql/patch.py

@@ -6,6 +6,7 @@
 from ddtrace.constants import SPAN_KIND
 from ddtrace.constants import SPAN_MEASURED_KEY
 from ddtrace.contrib import dbapi
+from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.schema import schematize_database_operation
 from ddtrace.internal.utils.wrappers import unwrap
@@ -16,12 +17,16 @@
 from ...ext import db
 from ...ext import net
 from ...internal.schema import schematize_service_name
+from ...propagation._database_monitoring import _DBM_Propagator
 from .. import trace_utils
 
 
 config._add(
     "aiomysql",
-    dict(_default_service=schematize_service_name("mysql")),
+    dict(
+        _default_service=schematize_service_name("mysql"),
+        _dbm_propagator=_DBM_Propagator(0, "query"),
+    ),
 )
 
 
@@ -43,7 +48,7 @@ async def patched_connect(connect_func, _, args, kwargs):
     tags = {}
     for tag, attr in CONN_ATTR_BY_TAG.items():
         if hasattr(conn, attr):
-            tags[tag] = getattr(conn, attr)
+            tags[tag] = trace_utils._convert_to_string(getattr(conn, attr, None))
     tags[db.SYSTEM] = "mysql"
 
     c = AIOTracedConnection(conn)
@@ -83,6 +88,11 @@ async def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
             # set analytics sample rate
             s.set_tag(ANALYTICS_SAMPLE_RATE_KEY, config.aiomysql.get_analytics_sample_rate())
 
+            # dispatch DBM
+            result = core.dispatch_with_results("aiomysql.execute", (config.aiomysql, s, args, kwargs)).result
+            if result:
+                s, args, kwargs = result.value
+
             try:
                 result = await method(*args, **kwargs)
                 return result

ddtrace/contrib/asyncpg/patch.py

@@ -4,6 +4,7 @@
 
 from ddtrace import Pin
 from ddtrace import config
+from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.vendor import wrapt
 
@@ -17,6 +18,7 @@
 from ...internal.schema import schematize_database_operation
 from ...internal.schema import schematize_service_name
 from ...internal.utils import get_argument_value
+from ...propagation._database_monitoring import _DBM_Propagator
 from ..trace_utils import ext_service
 from ..trace_utils import unwrap
 from ..trace_utils import wrap
@@ -37,6 +39,7 @@
     "asyncpg",
     dict(
         _default_service=schematize_service_name("postgres"),
+        _dbm_propagator=_DBM_Propagator(0, "query"),
     ),
 )
 
@@ -113,16 +116,21 @@ async def _traced_query(pin, method, query, args, kwargs):
 
         # set span.kind to the type of request being performed
         span.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
-
         span.set_tag(SPAN_MEASURED_KEY)
         span.set_tags(pin.tags)
+
+        # dispatch DBM
+        result = core.dispatch_with_results("asyncpg.execute", (config.asyncpg, method, span, args, kwargs)).result
+        if result:
+            span, args, kwargs = result.value
+
         return await method(*args, **kwargs)
 
 
 @with_traced_module
 async def _traced_protocol_execute(asyncpg, pin, func, instance, args, kwargs):
     state = get_argument_value(args, kwargs, 0, "state")  # type: Union[str, PreparedStatement]
-    query = state if isinstance(state, str) else state.query
+    query = state if isinstance(state, str) or isinstance(state, bytes) else state.query
     return await _traced_query(pin, func, query, args, kwargs)
 
 

ddtrace/contrib/dbapi/__init__.py

@@ -3,6 +3,7 @@
 """
 from ddtrace import config
 from ddtrace.appsec._iast._utils import _is_iast_enabled
+from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
 
 from ...appsec._constants import IAST_SPAN_TAGS
@@ -118,8 +119,14 @@ def _trace_method(self, method, name, resource, extra_tags, dbm_propagator, *arg
             if not isinstance(self, FetchTracedCursor):
                 s.set_tag(ANALYTICS_SAMPLE_RATE_KEY, self._self_config.get_analytics_sample_rate())
 
+            # dispatch DBM
             if dbm_propagator:
-                args, kwargs = dbm_propagator.inject(s, args, kwargs)
+                # this check is necessary to prevent fetch methods from trying to add dbm propagation
+                result = core.dispatch_with_results(
+                    f"{self._self_config.integration_name}.execute", (self._self_config, s, args, kwargs)
+                ).result
+                if result:
+                    s, args, kwargs = result.value
 
             try:
                 return method(*args, **kwargs)

ddtrace/contrib/dbapi_async/__init__.py

@@ -1,5 +1,6 @@
 from ddtrace import config
 from ddtrace.appsec._iast._utils import _is_iast_enabled
+from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
 
 from ...appsec._constants import IAST_SPAN_TAGS
@@ -88,8 +89,14 @@ async def _trace_method(self, method, name, resource, extra_tags, dbm_propagator
             if not isinstance(self, FetchTracedAsyncCursor):
                 s.set_tag(ANALYTICS_SAMPLE_RATE_KEY, self._self_config.get_analytics_sample_rate())
 
+            # dispatch DBM
             if dbm_propagator:
-                args, kwargs = dbm_propagator.inject(s, args, kwargs)
+                # this check is necessary to prevent fetch methods from trying to add dbm propagation
+                result = core.dispatch_with_results(
+                    f"{self._self_config.integration_name}.execute", [self._self_config, s, args, kwargs]
+                ).result
+                if result:
+                    s, args, kwargs = result.value
 
             try:
                 return await method(*args, **kwargs)

ddtrace/contrib/mysql/patch.py

@@ -12,6 +12,8 @@
 from ...internal.schema import schematize_database_operation
 from ...internal.schema import schematize_service_name
 from ...internal.utils.formats import asbool
+from ...propagation._database_monitoring import _DBM_Propagator
+from ..trace_utils import _convert_to_string
 
 
 config._add(
@@ -21,6 +23,7 @@
         _dbapi_span_name_prefix="mysql",
         _dbapi_span_operation_name=schematize_database_operation("mysql.query", database_provider="mysql"),
         trace_fetch_methods=asbool(os.getenv("DD_MYSQL_TRACE_FETCH_METHODS", default=False)),
+        _dbm_propagator=_DBM_Propagator(0, "query"),
     ),
 )
 
@@ -58,7 +61,9 @@ def _connect(func, instance, args, kwargs):
 
 
 def patch_conn(conn):
-    tags = {t: getattr(conn, a) for t, a in CONN_ATTR_BY_TAG.items() if getattr(conn, a, "") != ""}
+    tags = {
+        t: _convert_to_string(getattr(conn, a, None)) for t, a in CONN_ATTR_BY_TAG.items() if getattr(conn, a, "") != ""
+    }
     tags[db.SYSTEM] = "mysql"
     pin = Pin(tags=tags)
 

ddtrace/contrib/mysqldb/patch.py

@@ -19,6 +19,8 @@
 from ...internal.schema import schematize_service_name
 from ...internal.utils.formats import asbool
 from ...internal.utils.wrappers import unwrap as _u
+from ...propagation._database_monitoring import _DBM_Propagator
+from ..trace_utils import _convert_to_string
 
 
 config._add(
@@ -29,6 +31,7 @@
         _dbapi_span_operation_name=schematize_database_operation("mysql.query", database_provider="mysql"),
         trace_fetch_methods=asbool(os.getenv("DD_MYSQLDB_TRACE_FETCH_METHODS", default=False)),
         trace_connect=asbool(os.getenv("DD_MYSQLDB_TRACE_CONNECT", default=False)),
+        _dbm_propagator=_DBM_Propagator(0, "query"),
     ),
 )
 
@@ -99,7 +102,9 @@ def _connect(func, instance, args, kwargs):
 
 def patch_conn(conn, *args, **kwargs):
     tags = {
-        t: kwargs[k] if k in kwargs else args[p] for t, (k, p) in KWPOS_BY_TAG.items() if k in kwargs or len(args) > p
+        t: _convert_to_string(kwargs[k]) if k in kwargs else _convert_to_string(args[p])
+        for t, (k, p) in KWPOS_BY_TAG.items()
+        if k in kwargs or len(args) > p
     }
     tags[db.SYSTEM] = "mysql"
     tags[net.TARGET_PORT] = conn.port

ddtrace/contrib/pymysql/patch.py

@@ -12,6 +12,8 @@
 from ...internal.schema import schematize_database_operation
 from ...internal.schema import schematize_service_name
 from ...internal.utils.formats import asbool
+from ...propagation._database_monitoring import _DBM_Propagator
+from ..trace_utils import _convert_to_string
 
 
 config._add(
@@ -21,6 +23,7 @@
         _dbapi_span_name_prefix="pymysql",
         _dbapi_span_operation_name=schematize_database_operation("pymysql.query", database_provider="mysql"),
         trace_fetch_methods=asbool(os.getenv("DD_PYMYSQL_TRACE_FETCH_METHODS", default=False)),
+        _dbm_propagator=_DBM_Propagator(0, "query"),
     ),
 )
 
@@ -53,7 +56,7 @@ def _connect(func, instance, args, kwargs):
 
 
 def patch_conn(conn):
-    tags = {t: getattr(conn, a, "") for t, a in CONN_ATTR_BY_TAG.items()}
+    tags = {t: _convert_to_string(getattr(conn, a)) for t, a in CONN_ATTR_BY_TAG.items() if getattr(conn, a, "") != ""}
     tags[db.SYSTEM] = "mysql"
     pin = Pin(tags=tags)
 

ddtrace/contrib/trace_utils.py

@@ -23,6 +23,7 @@
 from ddtrace.ext import net
 from ddtrace.ext import user
 from ddtrace.internal import core
+from ddtrace.internal.compat import ensure_text
 from ddtrace.internal.compat import ip_is_global
 from ddtrace.internal.compat import parse
 from ddtrace.internal.logger import get_logger
@@ -671,3 +672,13 @@ def extract_netloc_and_query_info_from_url(url):
 
 class InterruptException(Exception):
     pass
+
+
+def _convert_to_string(attr):
+    # ensures attribute is converted to a string
+    if attr:
+        if isinstance(attr, int) or isinstance(attr, float):
+            return str(attr)
+        else:
+            return ensure_text(attr)
+    return attr

ddtrace/propagation/_database_monitoring.py

@@ -2,6 +2,7 @@
 from typing import Union  # noqa:F401
 
 import ddtrace
+from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.peer_service import PeerServiceConfig
 from ddtrace.vendor.sqlcommenter import generate_sql_comment as _generate_sql_comment
@@ -122,3 +123,28 @@ def _get_dbm_comment(self, db_span):
             # replace leading whitespace with trailing whitespace
             return sql_comment.strip() + " "
         return ""
+
+
+def handle_dbm_injection(int_config, span, args, kwargs):
+    dbm_propagator = getattr(int_config, "_dbm_propagator", None)
+    if dbm_propagator:
+        args, kwargs = dbm_propagator.inject(span, args, kwargs)
+
+    return span, args, kwargs
+
+
+def handle_dbm_injection_asyncpg(int_config, method, span, args, kwargs):
+    # bind_execute_many uses prepared statements which we want to avoid injection for
+    if method.__name__ != "bind_execute_many":
+        return handle_dbm_injection(int_config, span, args, kwargs)
+    return span, args, kwargs
+
+
+if dbm_config.propagation_mode in ["full", "service"]:
+    core.on("aiomysql.execute", handle_dbm_injection, "result")
+    core.on("asyncpg.execute", handle_dbm_injection_asyncpg, "result")
+    core.on("dbapi.execute", handle_dbm_injection, "result")
+    core.on("mysql.execute", handle_dbm_injection, "result")
+    core.on("mysqldb.execute", handle_dbm_injection, "result")
+    core.on("psycopg.execute", handle_dbm_injection, "result")
+    core.on("pymysql.execute", handle_dbm_injection, "result")

releasenotes/notes/add-dbm-for-mysql-and-postgres-integrations-8558939c45f0dbd7.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    aiomysql, asyncpg, mysql, mysqldb, pymysql: Add Database Monitoring (DBM) for remaining
+    mysql and postgres integrations lacking support.

tests/.suitespec.json

@@ -388,6 +388,9 @@
         "django": [
             "ddtrace/contrib/django/*"
         ],
+        "aiopg": [
+            "ddtrace/contrib/aiopg/*"
+        ],
         "pytest": [
             "ddtrace/contrib/pytest/*",
             "ddtrace/contrib/pytest_bdd/*",
@@ -731,14 +734,16 @@
             "@contrib",
             "@tracing",
             "@dbapi",
-            "@appsec"
+            "@appsec",
+            "tests/contrib/shared_tests.py"
         ],
         "dbapi_async": [
             "@bootstrap",
             "@core",
             "@contrib",
             "@tracing",
-            "@dbapi"
+            "@dbapi",
+            "tests/contrib/shared_tests.py"
         ],
         "asyncpg": [
             "@bootstrap",
@@ -747,7 +752,8 @@
             "@tracing",
             "@pg",
             "tests/contrib/asyncpg/*",
-            "tests/snapshots/tests.contrib.{suite}.*"
+            "tests/snapshots/tests.contrib.{suite}.*",
+            "tests/contrib/shared_tests.py"
         ],
         "aiohttp": [
             "@bootstrap",
@@ -1000,7 +1006,8 @@
             "@dbapi",
             "@mysql",
             "tests/contrib/mysqldb/*",
-            "tests/contrib/mysql/*"
+            "tests/contrib/mysql/*",
+            "tests/contrib/shared_tests.py"
         ],
         "pymysql": [
             "@bootstrap",
@@ -1010,7 +1017,8 @@
             "@dbapi",
             "@mysql",
             "tests/contrib/pymysql/*",
-            "tests/contrib/mysql/*"
+            "tests/contrib/mysql/*",
+            "tests/contrib/shared_tests.py"
         ],
         "pylibmc": [
             "@bootstrap",
@@ -1170,7 +1178,8 @@
             "@dbapi",
             "@pg",
             "tests/contrib/psycopg/*",
-            "tests/snapshots/tests.contrib.psycopg.*"
+            "tests/snapshots/tests.contrib.psycopg.*",
+            "tests/contrib/shared_tests.py"
         ],
         "psycopg2": [
             "@bootstrap",
@@ -1180,7 +1189,8 @@
             "@tracing",
             "@pg",
             "tests/contrib/psycopg2/*",
-            "tests/snapshots/tests.contrib.psycopg2.*"
+            "tests/snapshots/tests.contrib.psycopg2.*",
+            "tests/contrib/shared_tests.py"
         ],
         "aiobotocore": [
             "@bootstrap",
@@ -1198,7 +1208,8 @@
             "@dbapi",
             "@mysql",
             "tests/contrib/aiomysql/*",
-            "tests/snapshots/tests.contrib.{suite}.*"
+            "tests/snapshots/tests.contrib.{suite}.*",
+            "tests/contrib/shared_tests.py"
         ],
         "aiopg": [
             "@bootstrap",
@@ -1207,7 +1218,9 @@
             "@tracing",
             "@dbapi",
             "@pg",
-            "tests/contrib/aiopg/*"
+            "@aiopg",
+            "tests/contrib/aiopg/*",
+            "tests/contrib/shared_tests.py"
         ],
         "aredis": [
             "@bootstrap",