Prevent crash in SQL Server's JDBC when tracing execute methods with generated keys

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,

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