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 microsoft/mssql-jdbc#2729

Author: na-ji

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,

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}*/"
+  }
 }