Append comment if there's a pg plan hint

Author: edengorevoy

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;
+  }
 }

dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy

@@ -108,6 +108,7 @@ 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'*/"
   }
 
   def "test encode Sql Comment with peer service"() {