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..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( @@ -75,6 +72,7 @@ public static class StatementAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(value = 0, readOnly = false) String sql, + @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); @@ -127,6 +125,22 @@ 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; + + // 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 + // 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/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}*/" + } }