chore(dbm): generate sql comments for DBM (#4321)

## Description

Adds utils for generating DBM comments

## Checklist
- [ ] Title must conform to [conventional commit](https://github.com/conventional-changelog/commitlint/tree/master/%40commitlint/config-conventional).
- [ ] Add additional sections for `feat` and `fix` pull requests.
- [ ] Ensure tests are passing for affected code.
- [ ] [Library documentation](https://github.com/DataDog/dd-trace-py/tree/1.x/docs) and/or [Datadog's documentation site](https://github.com/DataDog/documentation/) is updated. Link to doc PR in description.





## Motivation


## Design 


## Testing strategy




## Relevant issue(s)


## Testing strategy




## Reviewer Checklist
- [ ] Title is accurate.
- [ ] Description motivates each change.
- [ ] No unnecessary changes were introduced in this PR.
- [ ] PR cannot be broken up into smaller PRs.
- [ ] Avoid breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes unless absolutely necessary.
- [ ] Tests provided or description of manual testing performed is included in the code or PR.
- [ ] Release note has been added for fixes and features, or else `changelog/no-changelog` label added.
- [ ] All relevant GitHub issues are correctly linked.
- [ ] Backports are identified and tagged with Mergifyio.
- [ ] Add to milestone.

Author: mabdinur

ddtrace/context.py

@@ -94,6 +94,15 @@ def sampling_priority(self, value):
                 return
             self._metrics[SAMPLING_PRIORITY_KEY] = value
 
+    @property
+    def _traceparent(self):
+        # type: () -> str
+        if self.trace_id is None or self.span_id is None:
+            return ""
+
+        sampled = 1 if self.sampling_priority and self.sampling_priority > 0 else 0
+        return "00-{:032x}-{:016x}-{:02x}".format(self.trace_id, self.span_id, sampled)
+
     @property
     def dd_origin(self):
         # type: () -> Optional[Text]

ddtrace/settings/database_monitoring.py

@@ -1,6 +1,22 @@
+from typing import TYPE_CHECKING
+
 from envier import En
 from envier import validators
 
+from ddtrace.vendor.sqlcommenter import generate_sql_comment as _generate_sql_comment
+
+from . import _config as dd_config
+
+
+if TYPE_CHECKING:
+    from ddtrace import Span
+
+DBM_PARENT_SERVICE_NAME_KEY = "ddps"
+DBM_DATABASE_SERVICE_NAME_KEY = "dddbs"
+DBM_ENVIRONMENT_KEY = "dde"
+DBM_VERSION_KEY = "ddpv"
+DBM_TRACE_PARENT_KEY = "traceparent"
+
 
 class DatabaseMonitoringConfig(En):
     __prefix__ = "dd_trace"
@@ -14,4 +30,23 @@ class DatabaseMonitoringConfig(En):
     )
 
 
-config = DatabaseMonitoringConfig()
+dbm_config = DatabaseMonitoringConfig()
+
+
+def _get_dbm_comment(db_span):
+    # type: (Span) -> str
+    if dbm_config.injection_mode == "disabled":
+        return ""
+
+    dbm_tags = {
+        DBM_PARENT_SERVICE_NAME_KEY: dd_config.service,
+        DBM_ENVIRONMENT_KEY: dd_config.env,
+        DBM_VERSION_KEY: dd_config.version,
+        DBM_DATABASE_SERVICE_NAME_KEY: db_span.service,
+    }
+
+    if dbm_config.injection_mode == "full":
+        # TODO: add _dd.dbm_trace_injected tag to db_span
+        dbm_tags[DBM_TRACE_PARENT_KEY] = db_span.context._traceparent
+
+    return _generate_sql_comment(**dbm_tags)

tests/internal/test_database_monitoring.py

@@ -7,7 +7,7 @@
 def test_injection_mode_configuration():
     # Ensure Database Monitoring support is disabled by default
     # TODO: Enable DBM support by default
-    assert database_monitoring.config.injection_mode == "disabled"
+    assert database_monitoring.dbm_config.injection_mode == "disabled"
 
     # Ensure service is a valid injection mode
     with override_env(dict(DD_TRACE_SQL_COMMENT_INJECTION_MODE="service")):
@@ -27,3 +27,55 @@ def test_injection_mode_configuration():
         excinfo.value.args[0] == "Invalid value for environment variable DD_TRACE_SQL_COMMENT_INJECTION_MODE: "
         "value must be one of ['disabled', 'full', 'service']"
     )
+
+
+@pytest.mark.subprocess(env=dict(DD_TRACE_SQL_COMMENT_INJECTION_MODE="disabled"))
+def test_get_dbm_comment_disabled_mode():
+    from ddtrace import tracer
+    from ddtrace.settings import database_monitoring
+
+    dbspan = tracer.trace("dbspan", service="orders-db")
+    sqlcomment = database_monitoring._get_dbm_comment(dbspan)
+    assert sqlcomment == ""
+
+
+@pytest.mark.subprocess(
+    env=dict(
+        DD_TRACE_SQL_COMMENT_INJECTION_MODE="service",
+        DD_SERVICE="orders-app",
+        DD_ENV="staging",
+        DD_VERSION="v7343437-d7ac743",
+    )
+)
+def test_get_dbm_comment_service_mode():
+    from ddtrace import tracer
+    from ddtrace.settings import database_monitoring
+
+    dbspan = tracer.trace("dbname", service="orders-db")
+
+    sqlcomment = database_monitoring._get_dbm_comment(dbspan)
+    # assert tags sqlcomment contains the correct value
+    assert sqlcomment == " /*dddbs='orders-db',dde='staging',ddps='orders-app',ddpv='v7343437-d7ac743'*/"
+
+
+@pytest.mark.subprocess(
+    env=dict(
+        DD_TRACE_SQL_COMMENT_INJECTION_MODE="full",
+        DD_SERVICE="orders-app",
+        DD_ENV="staging",
+        DD_VERSION="v7343437-d7ac743",
+    )
+)
+def test_get_dbm_comment_full_mode():
+    from ddtrace import tracer
+    from ddtrace.settings import database_monitoring
+
+    dbspan = tracer.trace("dbname", service="orders-db")
+
+    sqlcomment = database_monitoring._get_dbm_comment(dbspan)
+    # assert tags sqlcomment contains the correct value
+    assert (
+        sqlcomment
+        == " /*dddbs='orders-db',dde='staging',ddps='orders-app',ddpv='v7343437-d7ac743',traceparent='%s'*/"
+        % (dbspan.context._traceparent,)
+    )

tests/tracer/test_context.py

@@ -48,3 +48,30 @@ def test_eq(ctx1, ctx2):
 )
 def test_not_eq(ctx1, ctx2):
     assert ctx1 != ctx2
+
+
+def test_traceparent():
+    def validate_traceparent(context, sampled_expected):
+        version_hex, traceid_hex, spanid_hex, sampled_hex = context._traceparent.split("-")
+        assert version_hex == "00"
+
+        assert len(traceid_hex) == 32
+        assert traceid_hex == "{:032x}".format(context.trace_id)
+
+        assert len(spanid_hex) == 16
+        assert spanid_hex == "{:016x}".format(context.span_id)
+
+        assert len(sampled_hex) == 2
+        assert sampled_hex == sampled_expected
+
+    span = Span("span_a")
+    span.context.sampling_priority = -1
+    validate_traceparent(span.context, "00")
+
+    span = Span("span_b")
+    span.context.sampling_priority = 0
+    validate_traceparent(span.context, "00")
+
+    span = Span("span_c")
+    span.context.sampling_priority = 1
+    validate_traceparent(span.context, "01")