Add peer service tag in dbm sql commenter

dd-java-agent/instrumentation/jdbc/src/jmh/java/datadog/trace/instrumentation/jdbc/SQLCommenterBenchmark.java

@@ -33,6 +33,7 @@ public void testToComment() {
         dbService,
         hostname,
         dbName,
+        null,
         env,
         version,
         traceParent);

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java

@@ -1,6 +1,10 @@
 package datadog.trace.instrumentation.jdbc;
 
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
 import java.io.UnsupportedEncodingException;
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
@@ -15,6 +19,7 @@ public class SQLCommenter {
   private static final String DATABASE_SERVICE = encode("dddbs");
   private static final String DD_HOSTNAME = encode("ddh");
   private static final String DD_DB_NAME = encode("dddb");
+  private static final String DD_DB_INSTANCE = "ddprs";
   private static final String DD_ENV = encode("dde");
   private static final String DD_VERSION = encode("ddpv");
   private static final String TRACEPARENT = encode("traceparent");
@@ -94,13 +99,21 @@ public static String inject(
       }
     }
 
+    AgentSpan currSpan = activeSpan();
+    Object peerService = null;
+    if (currSpan != null) {
+      peerService = currSpan.getTag(Tags.PEER_SERVICE);
+    }
+
     final Config config = Config.get();
     final String parentService = config.getServiceName();
     final String env = config.getEnv();
     final String version = config.getVersion();
     final int commentSize = capacity(traceParent, parentService, dbService, env, version);
     StringBuilder sb = new StringBuilder(sql.length() + commentSize);
     boolean commentAdded = false;
+    String peerServiceString = peerService != null ? peerService.toString() : null;
+
     if (appendComment) {
       sb.append(sql);
       sb.append(SPACE);
@@ -113,6 +126,7 @@ public static String inject(
               dbService,
               hostname,
               dbName,
+              peerServiceString,
               env,
               version,
               traceParent);
@@ -127,9 +141,11 @@ public static String inject(
               dbService,
               hostname,
               dbName,
+              peerServiceString,
               env,
               version,
               traceParent);
+
       sb.append(CLOSE_COMMENT);
       sb.append(SPACE);
       sb.append(sql);
@@ -192,6 +208,7 @@ private static String encode(final String val) {
     return val;
   }
 
+  /*
   protected static boolean toComment(
       StringBuilder sb,
       final boolean injectTrace,
@@ -213,6 +230,34 @@ protected static boolean toComment(
       append(sb, TRACEPARENT, traceparent, sb.length() > emptySize);
     }
     return sb.length() > emptySize;
+  }*/
+
+  protected static boolean toComment(
+      StringBuilder sb,
+      final boolean injectTrace,
+      final String parentService,
+      final String dbService,
+      final String hostname,
+      final String dbName,
+      final String dbInstance,
+      final String env,
+      final String version,
+      final String traceparent) {
+    int emptySize = sb.length();
+
+    append(sb, PARENT_SERVICE, parentService, false);
+    append(sb, DATABASE_SERVICE, dbService, sb.length() > emptySize);
+    append(sb, DD_HOSTNAME, hostname, sb.length() > emptySize);
+    append(sb, DD_DB_NAME, dbName, sb.length() > emptySize);
+    if (dbInstance != null) {
+      append(sb, DD_DB_INSTANCE, dbInstance, sb.length() > emptySize);
+    }
+    append(sb, DD_ENV, env, sb.length() > emptySize);
+    append(sb, DD_VERSION, version, sb.length() > emptySize);
+    if (injectTrace) {
+      append(sb, TRACEPARENT, traceparent, sb.length() > emptySize);
+    }
+    return sb.length() > emptySize;
   }
 
   private static void append(StringBuilder sb, String key, String value, boolean prependComma) {

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

@@ -1,6 +1,12 @@
 import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.Tags
+
 import datadog.trace.instrumentation.jdbc.SQLCommenter
 
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+
 class SQLCommenterTest extends AgentTestRunner {
 
   def "test find first word"() {
@@ -105,4 +111,39 @@ class SQLCommenterTest extends AgentTestRunner {
     "/*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"
   }
+
+  def "test encode Sql Comment with peer service"() {
+    setup:
+    injectSysConfig("dd.service", ddService)
+    injectSysConfig("dd.env", ddEnv)
+    injectSysConfig("dd.version", ddVersion)
+
+    when:
+    String sqlWithComment = ""
+    runUnderTrace("testTrace"){
+      AgentSpan currSpan = AgentTracer.activeSpan()
+      currSpan.setTag(Tags.PEER_SERVICE, peerService)
+
+      if (injectTrace) {
+        sqlWithComment = SQLCommenter.inject(query, dbService, dbType, host, dbName, traceParent, true, appendComment)
+      } else {
+        if (appendComment) {
+          sqlWithComment = SQLCommenter.append(query, dbService, dbType, host, dbName)
+        } else {
+          sqlWithComment = SQLCommenter.prepend(query, dbService, dbType, host, dbName)
+        }
+      }
+    }
+    sqlWithComment == expected
+
+    then:
+    sqlWithComment == expected
+
+    where:
+    query                                                                                                         | ddService      | ddEnv  | dbService    | dbType     | host | dbName | ddVersion     | injectTrace | appendComment | traceParent                                               | peerService | expected
+    "SELECT * FROM foo"                                                                                           | "SqlCommenter" | "Test" | "my-service" | "mysql"    | "h"  | "n"    | "TestVersion" | true        | true          | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | ""          | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
+    "SELECT * FROM foo"                                                                                           | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h"  | "n"    | "TestVersion" | true        | true          | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | ""          | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
+    "SELECT * FROM foo"                                                                                           | "SqlCommenter" | "Test" | "my-service" | "mysql"    | "h"  | "n"    | "TestVersion" | true        | true          | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "testPeer" | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',ddprs='testPeer',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
+    "SELECT * FROM foo"                                                                                           | "SqlCommenter" | "Test" | "my-service" | "postgres" | "h"  | "n"    | "TestVersion" | true        | true          | "00-00000000000000007fffffffffffffff-000000024cb016ea-00" | "testPeer" | "SELECT * FROM foo /*ddps='SqlCommenter',dddbs='my-service',ddh='h',dddb='n',ddprs='testPeer',dde='Test',ddpv='TestVersion',traceparent='00-00000000000000007fffffffffffffff-000000024cb016ea-00'*/"
+  }
 }