Full mode for SQL Server

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

@@ -87,7 +87,7 @@ public Builder type(String type) {
       this.type = type;
       // Those DBs use the full text of the query including the comments as a cache key,
       // so we disable full propagation support for them to avoid destroying the cache.
-      if (type.equals("oracle") || type.equals("sqlserver")) this.fullPropagationSupport = false;
+      if (type.equals("oracle")) this.fullPropagationSupport = false;
       return this;
     }
 

dd-java-agent/instrumentation/jdbc/build.gradle

@@ -34,9 +34,11 @@ dependencies {
 
   testImplementation group: 'mysql', name: 'mysql-connector-java', version: '8.0.23'
   testImplementation group: 'org.postgresql', name: 'postgresql', version: '[9.4,42.2.18]'
+  testImplementation group: 'com.microsoft.sqlserver', name: 'mssql-jdbc', version: '10.2.0.jre8'
 
   testImplementation group: 'org.testcontainers', name:'mysql', version: libs.versions.testcontainers.get()
   testImplementation group: 'org.testcontainers', name:'postgresql', version: libs.versions.testcontainers.get()
+  testImplementation group: 'org.testcontainers', name:'mssqlserver', version: '1.19.8'
 
   testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter')
 

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

@@ -5,6 +5,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DATABASE_QUERY;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
+import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logMissingQueryInfo;
 import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logSQLException;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -16,6 +17,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
 import java.sql.Connection;
@@ -64,21 +66,31 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
         return null;
       }
       try {
-        Connection connection = statement.getConnection();
-        DBQueryInfo queryInfo =
+        final Connection connection = statement.getConnection();
+        final DBQueryInfo queryInfo =
             InstrumentationContext.get(Statement.class, DBQueryInfo.class).get(statement);
         if (null == queryInfo) {
           logMissingQueryInfo(statement);
           return null;
         }
-
-        final AgentSpan span = startSpan(DATABASE_QUERY);
-        DECORATE.afterStart(span);
-        DBInfo dbInfo =
+        final AgentSpan span;
+        final DBInfo dbInfo =
             JDBCDecorator.parseDBInfo(
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
+        final boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
+
+        if (INJECT_COMMENT && injectTraceContext && DECORATE.isSqlServer(dbInfo)) {
+          // The span ID is pre-determined so that we can reference it when setting the context
+          final long spanID = DECORATE.setContextInfo(connection, dbInfo);
+          // we then force that pre-determined span ID for the span covering the actual query
+          span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
+        } else {
+          span = startSpan(DATABASE_QUERY);
+        }
+        DECORATE.afterStart(span);
         DECORATE.onConnection(span, dbInfo);
         DECORATE.onPreparedStatement(span, queryInfo);
+
         return activateSpan(span);
       } catch (SQLException e) {
         logSQLException(e);

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

@@ -1,23 +1,30 @@
 package datadog.trace.instrumentation.jdbc;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_SCHEMA;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_WAREHOUSE;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
+import datadog.trace.api.DDTraceId;
 import datadog.trace.api.naming.SpanNaming;
 import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.DatabaseClientDecorator;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
 import datadog.trace.bootstrap.instrumentation.jdbc.JDBCConnectionUrlParser;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
+import java.sql.PreparedStatement;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.HashSet;
@@ -235,6 +242,73 @@ public String traceParent(AgentSpan span, int samplingPriority) {
     return sb.toString();
   }
 
+  public boolean isSqlServer(final DBInfo dbInfo) {
+    return "sqlserver".equals(dbInfo.getType());
+  }
+
+  /**
+   * Executes a `SET CONTEXT_INFO` statement on the DB with the active trace ID and the given span
+   * ID. This context will be "attached" to future queries on the same connection. See <a
+   * href="https://learn.microsoft.com/fr-fr/sql/t-sql/functions/context-info-transact-sql">MSSQL
+   * doc</a>. This is to be used where injecting trace and span in the comments with {@link
+   * SQLCommenter#inject} is not possible or convenient.
+   *
+   * <p>Upsides: still "visible" in sub-queries, does not bust caches based on full query text
+   * Downsides: takes time.
+   *
+   * @param connection The same connection as the one that will be used for the actual statement
+   * @param dbInfo dbInfo of the instrumented database
+   * @return spanID pre-created spanID
+   */
+  public long setContextInfo(Connection connection, DBInfo dbInfo) {
+    final byte VERSION = 0;
+    final long spanID = Config.get().getIdGenerationStrategy().generateSpanId();
+    AgentSpan instrumentationSpan =
+        AgentTracer.get().buildSpan("set context_info").withTag("dd.instrumentation", true).start();
+    DECORATE.afterStart(instrumentationSpan);
+    DECORATE.onConnection(instrumentationSpan, dbInfo);
+    PreparedStatement instrumentationStatement = null;
+    try (AgentScope scope = activateSpan(instrumentationSpan)) {
+      final byte samplingDecision =
+          (byte) (instrumentationSpan.forceSamplingDecision() > 0 ? 1 : 0);
+      final byte versionAndSamplingDecision =
+          (byte) ((VERSION << 4) & 0b11110000 | samplingDecision & 0b00000001);
+
+      ByteBuffer byteBuffer = ByteBuffer.allocate(1 + 3 * Long.BYTES);
+      byteBuffer.order(ByteOrder.BIG_ENDIAN);
+
+      byteBuffer.put(versionAndSamplingDecision);
+      byteBuffer.putLong(spanID);
+      final DDTraceId traceId = instrumentationSpan.getTraceId();
+      byteBuffer.putLong(traceId.toHighOrderLong());
+      byteBuffer.putLong(traceId.toLong());
+      final byte[] contextInfo = byteBuffer.array();
+
+      String instrumentationSql = "set context_info ?";
+      instrumentationStatement = connection.prepareStatement(instrumentationSql);
+      instrumentationStatement.setBytes(1, contextInfo);
+      DECORATE.onStatement(instrumentationSpan, instrumentationSql);
+      instrumentationStatement.execute();
+    } catch (Exception e) {
+      log.debug(
+          "Failed to set extra DBM data in context info for trace {}. "
+              + "To disable this behavior, set DBM_PROPAGATION_MODE to 'service' mode. "
+              + "See https://docs.datadoghq.com/database_monitoring/connect_dbm_and_apm/ for more info.{}",
+          instrumentationSpan.getTraceId().toHexString(),
+          e);
+      DECORATE.onError(instrumentationSpan, e);
+    } finally {
+      if (instrumentationStatement != null) {
+        try {
+          instrumentationStatement.close();
+        } catch (Exception e) {
+        }
+      }
+      instrumentationSpan.finish();
+    }
+    return spanID;
+  }
+
   @Override
   protected void postProcessServiceAndOperationName(
       AgentSpan span, DatabaseClientDecorator.NamingEntry namingEntry) {

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

@@ -21,6 +21,7 @@
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
 import java.sql.Connection;
 import java.sql.SQLException;
@@ -84,18 +85,29 @@ public static AgentScope onEnter(
       }
       try {
         final Connection connection = statement.getConnection();
-        final AgentSpan span = startSpan(DATABASE_QUERY);
-        DECORATE.afterStart(span);
         final DBInfo dbInfo =
             JDBCDecorator.parseDBInfo(
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
+        boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
+        final AgentSpan span;
+        final boolean isSqlServer = DECORATE.isSqlServer(dbInfo);
+
+        if (isSqlServer && INJECT_COMMENT && injectTraceContext) {
+          // The span ID is pre-determined so that we can reference it when setting the context
+          final long spanID = DECORATE.setContextInfo(connection, dbInfo);
+          // we then force that pre-determined span ID for the span covering the actual query
+          span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
+        } else {
+          span = startSpan(DATABASE_QUERY);
+        }
+
+        DECORATE.afterStart(span);
         DECORATE.onConnection(span, dbInfo);
         final String copy = sql;
         if (span != null && INJECT_COMMENT) {
           String traceParent = null;
 
-          boolean injectTraceContext = DECORATE.shouldInjectTraceContext(dbInfo);
-          if (injectTraceContext) {
+          if (injectTraceContext && !isSqlServer) {
             Integer priority = span.forceSamplingDecision();
             if (priority != null) {
               traceParent = DECORATE.traceParent(span, priority);

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

@@ -21,7 +21,7 @@ abstract class JDBCDecoratorTest extends AgentTestRunner {
     where:
     dbType       | expectedByType
     "oracle"     | false
-    "sqlserver"  | false
+    "sqlserver"  | true
     "mysql"      | true
     "postgresql" | true
   }