Add option to JDBC instrumentation to always append DBM comment

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java

@@ -126,7 +126,9 @@ public static String onEnter(
         dbService =
             traceConfig(activeSpan()).getServiceMapping().getOrDefault(dbService, dbService);
       }
-      boolean append = "sqlserver".equals(dbInfo.getType());
+
+      boolean append =
+          DECORATE.DBM_ALWAYS_APPEND_SQL_COMMENT || "sqlserver".equals(dbInfo.getType());
       sql =
           SQLCommenter.inject(
               sql, dbService, dbInfo.getType(), dbInfo.getHost(), dbInfo.getDb(), null, append);

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java

@@ -59,6 +59,8 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
       DBM_PROPAGATION_MODE.equals(DBM_PROPAGATION_MODE_FULL);
   public static final boolean DBM_TRACE_PREPARED_STATEMENTS =
       Config.get().isDbmTracePreparedStatements();
+  public static final boolean DBM_ALWAYS_APPEND_SQL_COMMENT =
+      Config.get().isDbmAlwaysAppendSqlComment();
 
   private volatile boolean warnedAboutDBMPropagationMode = false; // to log a warning only once
 

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java

@@ -127,14 +127,15 @@ public static AgentScope onEnter(
           final boolean injectTraceInComment = injectTraceContext && !isSqlServer && !isOracle;
 
           // prepend mode will prepend the SQL comment to the raw sql query
-          boolean appendComment = false;
+          boolean appendComment = DECORATE.DBM_ALWAYS_APPEND_SQL_COMMENT;
 
           // There is a bug in the SQL Server JDBC driver that prevents
           // the generated keys from being returned when the
           // SQL comment is prepended to the SQL query.
           // We only append in this case to avoid the comment from being truncated.
           // @see https://github.com/microsoft/mssql-jdbc/issues/2729
           if (isSqlServer
+              && !appendComment
               && args.length == 2
               && args[1] instanceof Integer
               && (Integer) args[1] == Statement.RETURN_GENERATED_KEYS) {

dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy

@@ -5,8 +5,7 @@ import test.TestConnection
 import test.TestPreparedStatement
 import test.TestStatement
 
-class DBMInjectionForkedTest extends InstrumentationSpecification {
-
+abstract class InjectionTest extends InstrumentationSpecification {
   @Override
   void configurePreAgent() {
     super.configurePreAgent()
@@ -21,8 +20,9 @@ class DBMInjectionForkedTest extends InstrumentationSpecification {
 
   static query = "SELECT 1"
   static serviceInjection = "ddps='my_service_name',dddbs='remapped_testdb',ddh='localhost'"
-  static fullInjection = serviceInjection + ",traceparent='00-00000000000000000000000000000004-0000000000000003-01'"
+}
 
+class DBMInjectionForkedTest extends InjectionTest {
   def "prepared stmt"() {
     setup:
     def connection = new TestConnection(false)
@@ -45,6 +45,35 @@ class DBMInjectionForkedTest extends InstrumentationSpecification {
     statement.executeQuery(query)
 
     then:
-    assert statement.sql == "/*${fullInjection}*/ ${query}"
+    assert statement.sql == "/*${serviceInjection},traceparent='00-00000000000000000000000000000004-0000000000000003-01'*/ ${query}"
+  }
+}
+
+class DBMAppendInjectionForkedTest extends InjectionTest {
+  def "append comment on prepared stmt"() {
+    setup:
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT, "true")
+    def connection = new TestConnection(false)
+
+    when:
+    def statement = connection.prepareStatement(query) as TestPreparedStatement
+    statement.execute()
+
+    then:
+    // even in full propagation mode, we cannot inject trace info in prepared statements
+    assert statement.sql == "${query} /*${serviceInjection}*/"
+  }
+
+  def "append comment on single query"() {
+    setup:
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT, "true")
+    def connection = new TestConnection(false)
+
+    when:
+    def statement = connection.createStatement() as TestStatement
+    statement.executeQuery(query)
+
+    then:
+    assert statement.sql == "${query} /*${serviceInjection},traceparent='00-00000000000000000000000000000004-0000000000000003-01'*/"
   }
 }

dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLServerInjectionForkedTest.groovy

@@ -1,13 +1,10 @@
-
-
 import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.api.config.TraceInstrumentationConfig
 import test.TestConnection
 import test.TestDatabaseMetaData
 import test.TestStatement
 
-class SQLServerInjectionForkedTest extends InstrumentationSpecification {
-
+abstract class InjectionForkedTest extends InstrumentationSpecification {
   @Override
   void configurePreAgent() {
     super.configurePreAgent()
@@ -18,7 +15,9 @@ class SQLServerInjectionForkedTest extends InstrumentationSpecification {
 
   static query = "SELECT 1"
   static serviceInjection = "ddps='my_service_name',dddbs='sqlserver',ddh='localhost',dddb='testdb'"
+}
 
+class SQLServerInjectionForkedTest extends InjectionForkedTest {
   def "SQL Server no trace injection with full propagation mode"() {
     setup:
     def connection = new TestConnection(false)
@@ -37,7 +36,7 @@ class SQLServerInjectionForkedTest extends InstrumentationSpecification {
     assert !statement.sql.contains("traceparent")
   }
 
-  def "SQL Server apend comment when getting generated keys"() {
+  def "SQL Server append comment when getting generated keys"() {
     setup:
     def connection = new TestConnection(false)
     def metadata = new TestDatabaseMetaData()
@@ -52,3 +51,21 @@ class SQLServerInjectionForkedTest extends InstrumentationSpecification {
     assert statement.sql == "${query} /*${serviceInjection}*/"
   }
 }
+
+class SQLServerAppendInjectionForkedTest extends InjectionForkedTest {
+  def "SQL Server append comment when configured to do so"() {
+    setup:
+    injectSysConfig(TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT, "true")
+    def connection = new TestConnection(false)
+    def metadata = new TestDatabaseMetaData()
+    metadata.setURL("jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=testdb;")
+    connection.setMetaData(metadata)
+
+    when:
+    def statement = connection.createStatement() as TestStatement
+    statement.executeUpdate(query)
+
+    then:
+    assert statement.sql == "${query} /*${serviceInjection}*/"
+  }
+}

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -72,6 +72,7 @@ public final class ConfigDefaults {
   static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_HOST = false;
   static final String DEFAULT_DB_DBM_PROPAGATION_MODE_MODE = "disabled";
   static final boolean DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS = false;
+  static final boolean DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT = false;
   // Default value is set to 0, it disables the latency trace interceptor
   static final int DEFAULT_TRACE_KEEP_LATENCY_THRESHOLD_MS = 0;
   static final int DEFAULT_SCOPE_DEPTH_LIMIT = 100;

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -69,6 +69,7 @@ public final class TraceInstrumentationConfig {
   public static final String DB_DBM_INJECT_SQL_BASEHASH = "dbm.inject.sql.basehash";
   public static final String DB_DBM_PROPAGATION_MODE_MODE = "dbm.propagation.mode";
   public static final String DB_DBM_TRACE_PREPARED_STATEMENTS = "dbm.trace_prepared_statements";
+  public static final String DB_DBM_ALWAYS_APPEND_SQL_COMMENT = "dbm.always_append_sql_comment";
 
   public static final String JDBC_CONNECTION_CLASS_NAME = "trace.jdbc.connection.class.name";
 

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -55,6 +55,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_CLIENT_HOST_SPLIT_BY_HOST;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_PROPAGATION_MODE_MODE;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DEBUGGER_EXCEPTION_CAPTURE_INTERMEDIATE_SPANS_ENABLED;
@@ -508,6 +509,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_HOST;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX;
+import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_INJECT_SQL_BASEHASH;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_TRACE_PREPARED_STATEMENTS;
@@ -1076,6 +1078,7 @@ public static String getHostName() {
   private final boolean dbmInjectSqlBaseHash;
   private final String dbmPropagationMode;
   private final boolean dbmTracePreparedStatements;
+  private final boolean dbmAlwaysAppendSqlComment;
 
   private final boolean dynamicInstrumentationEnabled;
   private final String dynamicInstrumentationSnapshotUrl;
@@ -1639,6 +1642,10 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
         configProvider.getBoolean(
             DB_DBM_TRACE_PREPARED_STATEMENTS, DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS);
 
+    dbmAlwaysAppendSqlComment =
+        configProvider.getBoolean(
+            DB_DBM_ALWAYS_APPEND_SQL_COMMENT, DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT);
+
     dbmInjectSqlBaseHash = configProvider.getBoolean(DB_DBM_INJECT_SQL_BASEHASH, false);
 
     splitByTags = tryMakeImmutableSet(configProvider.getList(SPLIT_BY_TAGS));
@@ -5042,6 +5049,10 @@ public boolean isDbmTracePreparedStatements() {
     return dbmTracePreparedStatements;
   }
 
+  public boolean isDbmAlwaysAppendSqlComment() {
+    return dbmAlwaysAppendSqlComment;
+  }
+
   public String getDbmPropagationMode() {
     return dbmPropagationMode;
   }

metadata/supported-configurations.json

@@ -139,6 +139,7 @@
     "DD_DATA_STREAMS_BUCKET_DURATION_SECONDS": ["A"],
     "DD_DATA_STREAMS_ENABLED": ["A"],
     "DD_DBM_INJECT_SQL_BASEHASH": ["A"],
+    "DD_DBM_ALWAYS_APPEND_SQL_COMMENT": ["A"],
     "DD_DBM_PROPAGATION_MODE": ["A"],
     "DD_DBM_TRACE_PREPARED_STATEMENTS": ["A"],
     "DD_DISTRIBUTED_DEBUGGER_ENABLED": ["A"],