Fix

DataDog · Feb 28, 2024 · 1f9d4e7 · 1f9d4e7
1 parent 73a7404
commit 1f9d4e7
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 9 deletions.
diff --git a/.../main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java b/.../main/java/datadog/trace/instrumentation/jdbc/DBMCompatibleConnectionInstrumentation.java
@@ -100,9 +100,9 @@ public static String onEnter(
             JDBCDecorator.parseDBInfo(
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
         if (dbInfo.getType().equals("sqlserver")) {
-          sql = SQLCommenter.append(sql, DECORATE.getDbService(dbInfo));
+          sql = SQLCommenter.append(sql, DECORATE.getDbService(dbInfo), dbInfo.getType());
         } else {
-          sql = SQLCommenter.prepend(sql, DECORATE.getDbService(dbInfo));
+          sql = SQLCommenter.prepend(sql, DECORATE.getDbService(dbInfo), dbInfo.getType());
         }
         return inputSql;
       }

diff --git a/...jdbc/src/main/java/datadog/trace/instrumentation/jdbc/MySQLConnectionInstrumentation.java b/...jdbc/src/main/java/datadog/trace/instrumentation/jdbc/MySQLConnectionInstrumentation.java
@@ -98,7 +98,7 @@ public static String onEnter(
         final DBInfo dbInfo =
             JDBCDecorator.parseDBInfo(
                 connection, InstrumentationContext.get(Connection.class, DBInfo.class));
-        sql = SQLCommenter.prepend(sql, DECORATE.getDbService(dbInfo));
+        sql = SQLCommenter.prepend(sql, DECORATE.getDbService(dbInfo), dbInfo.getType());
         return inputSql;
       }
       return sql;

diff --git a/...t/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java b/...t/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java
@@ -24,17 +24,18 @@ public class SQLCommenter {
   private static final String CLOSE_COMMENT = "*/";
   private static final int INITIAL_CAPACITY = computeInitialCapacity();
 
-  public static String append(final String sql, final String dbService) {
-    return inject(sql, dbService, null, false, true);
+  public static String append(final String sql, final String dbService, final String dbType) {
+    return inject(sql, dbService, dbType, null, false, true);
   }
 
-  public static String prepend(final String sql, final String dbService) {
-    return inject(sql, dbService, null, false, false);
+  public static String prepend(final String sql, final String dbService, final String dbType) {
+    return inject(sql, dbService, dbType, null, false, false);
   }
 
   public static String inject(
       final String sql,
       final String dbService,
+      final String dbType,
       final String traceParent,
       final boolean injectTrace,
       final boolean appendComment) {
@@ -45,7 +46,7 @@ public static String inject(
       return sql;
     }
 
-    if (dbService.startsWith("postgres")) {
+    if (dbType != null && dbType.startsWith("postgres")) {
       //      The Postgres JDBC parser doesn't allow SQL comments anywhere in a callable statement
       //
       // https://github.com/pgjdbc/pgjdbc/blob/master/pgjdbc/src/main/java/org/postgresql/core/Parser.java#L1038

diff --git a/...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java b/...ation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java
@@ -103,7 +103,12 @@ public static AgentScope onEnter(
           }
           sql =
               SQLCommenter.inject(
-                  sql, span.getServiceName(), traceParent, injectTraceContext, appendComment);
+                  sql,
+                  span.getServiceName(),
+                  dbInfo.getType(),
+                  traceParent,
+                  injectTraceContext,
+                  appendComment);
         }
         DECORATE.onStatement(span, DBQueryInfo.ofStatement(copy));
         return activateSpan(span);