Do not prepend DBM <> APM trace comment in SQLCommenter if there is a pg plan hint (DataDog#8864)

edengorevoy · web-flow · commit d68f8f1e3037 · 2025-05-22T12:08:22.000-04:00
* Append comment if there's a pg plan hint

* Another test for comment placement
diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java
@@ -80,9 +80,11 @@ public static String inject(
     if (dbType != null) {
       final String firstWord = getFirstWord(sql);
 
-      // The Postgres JDBC parser doesn't allow SQL comments anywhere in a JDBC callable statements
+      // The Postgres JDBC parser doesn't allow SQL comments anywhere in a JDBC
+      // callable statements
       // https://github.com/pgjdbc/pgjdbc/blob/master/pgjdbc/src/main/java/org/postgresql/core/Parser.java#L1038
-      // TODO: Could we inject the comment after the JDBC has been converted to standard SQL?
+      // TODO: Could we inject the comment after the JDBC has been converted to
+      // standard SQL?
       if (firstWord.startsWith("{") && dbType.startsWith("postgres")) {
         return sql;
       }
@@ -92,11 +94,17 @@ public static String inject(
         appendComment = true;
       }
 
-      // Both Postgres and MySQL are unhappy with anything before CALL in a stored procedure
+      // Both Postgres and MySQL are unhappy with anything before CALL in a stored
+      // procedure
       // invocation but they seem ok with it after so we force append mode
       if (firstWord.equalsIgnoreCase("call")) {
         appendComment = true;
       }
+
+      // Append the comment in the case of a pg_hint_plan extension
+      if (dbType.startsWith("postgres") && containsPgHint(sql)) {
+        appendComment = true;
+      }
     }
 
     AgentSpan currSpan = activeSpan();
@@ -295,4 +303,10 @@ private static int computeInitialCapacity() {
             + CLOSE_COMMENT.length(); // two quotes, one equals & one comma * 5 + \* */
     return tagKeysLen + extraCharsLen;
   }
+
+  // pg_hint_plan extension works by checking the first block comment
+  // we'll have to append the traced comment if there is a pghint
+  private static boolean containsPgHint(String sql) {
+    return sql.indexOf("/*+") > 0;
+  }
 }
diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy
@@ -108,6 +108,8 @@ class SQLCommenterTest extends AgentTestRunner {
     "/*traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-01'*/ SELECT * FROM foo"                 | "SqlCommenter" | "Test" | "my-service" | "mysql"    | "h"  | "n"    | "TestVersion" | false       | false         | null                                                      | "/*traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-01'*/ SELECT * FROM foo"
     "/*customer-comment*/ SELECT * FROM foo"                                                                      | "SqlCommenter" | "Test" | "my-service" | "mysql"    | "h"  | "n"    | "TestVersion" | false       | false         | null                                                      | "/*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion'*/ /*customer-comment*/ SELECT * FROM foo"
     "/*traceparent"                                                                                               | "SqlCommenter" | "Test" | "my-service" | "mysql"    | "h"  | "n"    | "TestVersion" | false       | false         | null                                                      | "/*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion'*/ /*traceparent"
+    "SELECT /*+ SeqScan(foo) */ * FROM foo"                                                                       | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h"  | "n" | "TestVersion"    | true        | true          | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" |  "SELECT /*+ SeqScan(foo) */ * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
+    "/*+ SeqScan(foo) */ SELECT * FROM foo"                                                                       | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h"  | "n" | "TestVersion"    | true        | true          | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" |  "/*+ SeqScan(foo) */ SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
   }
 
   def "test encode Sql Comment with peer service"() {

Original file line number	Diff line number	Diff line change
`@@ -108,6 +108,8 @@ class SQLCommenterTest extends AgentTestRunner {`
`108`	`108`	`"/traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-01'/ SELECT * FROM foo" \| "SqlCommenter" \| "Test" \| "my-service" \| "mysql" \| "h" \| "n" \| "TestVersion" \| false \| false \| null \| "/traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-01'/ SELECT * FROM foo"`
`109`	`109`	`"/customer-comment/ SELECT * FROM foo" \| "SqlCommenter" \| "Test" \| "my-service" \| "mysql" \| "h" \| "n" \| "TestVersion" \| false \| false \| null \| "/ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion'/ /customer-comment/ SELECT * FROM foo"`
`110`	`110`	`"/traceparent" \| "SqlCommenter" \| "Test" \| "my-service" \| "mysql" \| "h" \| "n" \| "TestVersion" \| false \| false \| null \| "/ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion'/ /traceparent"`
	`111`	`+ "SELECT /+ SeqScan(foo) / * FROM foo" \| "SqlCommenter" \| "Test" \| "my-service" \| "postgres" \| "h" \| "n" \| "TestVersion" \| true \| true \| "00-00000000000000007fffffffffffffff-000000024cb016ea-00" \| "SELECT /+ SeqScan(foo) / * FROM foo /ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'/"`
	`112`	`+ "/+ SeqScan(foo) / SELECT * FROM foo" \| "SqlCommenter" \| "Test" \| "my-service" \| "postgres" \| "h" \| "n" \| "TestVersion" \| true \| true \| "00-00000000000000007fffffffffffffff-000000024cb016ea-00" \| "/+ SeqScan(foo) / SELECT * FROM foo /ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'/"`
`111`	`113`	`}`
`112`	`114`
`113`	`115`	`def "test encode Sql Comment with peer service"() {`