From 69211ede4f061232e157da642487e0fa93f69fa4 Mon Sep 17 00:00:00 2001 From: Naji Astier Date: Wed, 6 Aug 2025 12:00:54 +0200 Subject: [PATCH 1/4] Prevent crash in SQL Server JDBC when tracing execute methods with generated keys When tracing the JDBC statement execute methods, we prepend a comment to the SQL query used for DBM trace propagation. However, 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. I decided to only append the comment in this specific case to avoid the comment from being truncated. @see https://github.com/microsoft/mssql-jdbc/issues/2729 --- .../jdbc/StatementInstrumentation.java | 17 +++++++++++++++++ .../test/groovy/DBMInjectionForkedTest.groovy | 12 ++++++++++++ 2 files changed, 29 insertions(+) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java index 654d98dbed4..683b34317ce 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java @@ -10,6 +10,7 @@ import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE; import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT; import static java.util.Collections.singletonMap; +import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -75,6 +76,7 @@ public static class StatementAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(value = 0, readOnly = false) String sql, + @Advice.AllArguments(typing = DYNAMIC) Object[] args, @Advice.This final Statement statement) { // TODO consider matching known non-wrapper implementations to avoid this check final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class); @@ -127,6 +129,21 @@ public static AgentScope onEnter( // context_info and v$session.action respectively. // we should not also inject it into SQL comments to avoid duplication final boolean injectTraceInComment = injectTraceContext && !isSqlServer && !isOracle; + + boolean appendComment = StatementInstrumentation.appendComment; + + // 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 + && args.length == 2 + && args[1] instanceof Integer + && (Integer) args[1] == Statement.RETURN_GENERATED_KEYS) { + appendComment = true; + } + sql = SQLCommenter.inject( sql, diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy index 12ecd5e4b85..6f1e8066f8a 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy @@ -47,4 +47,16 @@ class DBMInjectionForkedTest extends AgentTestRunner { then: assert statement.sql == "/*${fullInjection}*/ ${query}" } + + def "single query with generated keys"() { + setup: + def connection = new TestConnection(false) + + when: + def statement = connection.createStatement() as TestStatement + statement.executeUpdate(query, 1) + + then: + assert statement.sql == "${query} /*${fullInjection}*/" + } } From 94d3981ef5f2d66a269219211e97068e310f842c Mon Sep 17 00:00:00 2001 From: Naji Astier Date: Thu, 14 Aug 2025 10:44:06 +0200 Subject: [PATCH 2/4] remove dynamic typing when packing all the JDBC arguments --- .../trace/instrumentation/jdbc/StatementInstrumentation.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java index 683b34317ce..fc4d428d017 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java @@ -10,7 +10,6 @@ import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE; import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT; import static java.util.Collections.singletonMap; -import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; @@ -76,7 +75,7 @@ public static class StatementAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(value = 0, readOnly = false) String sql, - @Advice.AllArguments(typing = DYNAMIC) Object[] args, + @Advice.AllArguments() Object[] args, @Advice.This final Statement statement) { // TODO consider matching known non-wrapper implementations to avoid this check final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class); From e84d9ebef334ddd80ed9d3b9f1ee084a008009d6 Mon Sep 17 00:00:00 2001 From: Naji Astier Date: Thu, 14 Aug 2025 10:44:43 +0200 Subject: [PATCH 3/4] move JDBC test on injection with generated keys to SQL Server file --- .../src/test/groovy/DBMInjectionForkedTest.groovy | 12 ------------ .../groovy/SQLServerInjectionForkedTest.groovy | 15 +++++++++++++++ 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy index 6f1e8066f8a..12ecd5e4b85 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/DBMInjectionForkedTest.groovy @@ -47,16 +47,4 @@ class DBMInjectionForkedTest extends AgentTestRunner { then: assert statement.sql == "/*${fullInjection}*/ ${query}" } - - def "single query with generated keys"() { - setup: - def connection = new TestConnection(false) - - when: - def statement = connection.createStatement() as TestStatement - statement.executeUpdate(query, 1) - - then: - assert statement.sql == "${query} /*${fullInjection}*/" - } } diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLServerInjectionForkedTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLServerInjectionForkedTest.groovy index 74829719b87..e2cf09b8ddd 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLServerInjectionForkedTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLServerInjectionForkedTest.groovy @@ -36,4 +36,19 @@ class SQLServerInjectionForkedTest extends AgentTestRunner { // Verify that the SQL does NOT contain traceparent assert !statement.sql.contains("traceparent") } + + def "SQL Server apend comment when getting generated keys"() { + setup: + 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, 1) + + then: + assert statement.sql == "${query} /*${serviceInjection}*/" + } } From 3a324e0a0ee4faebd1c5ce286ef85db6115be4b1 Mon Sep 17 00:00:00 2001 From: Naji Astier Date: Fri, 15 Aug 2025 16:29:59 +0200 Subject: [PATCH 4/4] remove useless constant in StatementAdvice --- .../instrumentation/jdbc/StatementInstrumentation.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java index fc4d428d017..23d657516dc 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java @@ -61,9 +61,6 @@ public String[] helperClassNames() { return new String[] {packageName + ".JDBCDecorator", packageName + ".SQLCommenter"}; } - // prepend mode will prepend the SQL comment to the raw sql query - private static final boolean appendComment = false; - @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( @@ -129,7 +126,8 @@ public static AgentScope onEnter( // we should not also inject it into SQL comments to avoid duplication final boolean injectTraceInComment = injectTraceContext && !isSqlServer && !isOracle; - boolean appendComment = StatementInstrumentation.appendComment; + // prepend mode will prepend the SQL comment to the raw sql query + boolean appendComment = false; // There is a bug in the SQL Server JDBC driver that prevents // the generated keys from being returned when the