From 6ec84551aa297df3f47a9e23f3cb4ceec1e438cc Mon Sep 17 00:00:00 2001 From: Jordan Wong <61242306+jordan-wong@users.noreply.github.com> Date: Thu, 30 Jan 2025 10:40:02 -0500 Subject: [PATCH] Add peer service tag in dbm sql commenter (#7913) * add peer.service to the SQLCommenter when explicity specified, add tests for it * Remove comments, cleanup * add working version with unit tests passing * remove comments and dev print statements * remove accidentally added extra param * add null check for ActiveSpan * refactor toComment behavior * fix small bug * clean up PR, rename variables and refactor test logic for readability --- .../jdbc/SQLCommenterBenchmark.java | 1 + .../instrumentation/jdbc/SQLCommenter.java | 21 ++++++++ .../src/test/groovy/SQLCommenterTest.groovy | 52 ++++++++++++++++--- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/jmh/java/datadog/trace/instrumentation/jdbc/SQLCommenterBenchmark.java b/dd-java-agent/instrumentation/jdbc/src/jmh/java/datadog/trace/instrumentation/jdbc/SQLCommenterBenchmark.java index e2a5b195416..195b221d2bc 100644 --- a/dd-java-agent/instrumentation/jdbc/src/jmh/java/datadog/trace/instrumentation/jdbc/SQLCommenterBenchmark.java +++ b/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); 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 index 97925d561c8..979a08e7232 100644 --- 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 @@ -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_PEER_SERVICE = "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,6 +99,12 @@ public static String inject( } } + AgentSpan currSpan = activeSpan(); + Object peerServiceObj = null; + if (currSpan != null) { + peerServiceObj = currSpan.getTag(Tags.PEER_SERVICE); + } + final Config config = Config.get(); final String parentService = config.getServiceName(); final String env = config.getEnv(); @@ -101,6 +112,8 @@ public static String inject( final int commentSize = capacity(traceParent, parentService, dbService, env, version); StringBuilder sb = new StringBuilder(sql.length() + commentSize); boolean commentAdded = false; + String peerService = peerServiceObj != null ? peerServiceObj.toString() : null; + if (appendComment) { sb.append(sql); sb.append(SPACE); @@ -113,6 +126,7 @@ public static String inject( dbService, hostname, dbName, + peerService, env, version, traceParent); @@ -127,9 +141,11 @@ public static String inject( dbService, hostname, dbName, + peerService, env, version, traceParent); + sb.append(CLOSE_COMMENT); sb.append(SPACE); sb.append(sql); @@ -199,14 +215,19 @@ protected static boolean toComment( final String dbService, final String hostname, final String dbName, + final String peerService, 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 (peerService != null) { + append(sb, DD_PEER_SERVICE, peerService, sb.length() > emptySize); + } append(sb, DD_ENV, env, sb.length() > emptySize); append(sb, DD_VERSION, version, sb.length() > emptySize); if (injectTrace) { diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy index 5af1054460d..a91f4b1d386 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy +++ b/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"() { @@ -35,15 +41,13 @@ class SQLCommenterTest extends AgentTestRunner { String sqlWithComment = "" 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) - } + } 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 @@ -105,4 +109,38 @@ 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) + } + } + + 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'*/" + } }