Skip to content

Commit

Permalink
Add peer service tag in dbm sql commenter (#7913)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@jordan-wong
jordan-wong authored Jan 30, 2025
1 parent a6b67c3 commit 6ec8455
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 7 deletions.
1 change: 1 addition & 0 deletions ...mentation/jdbc/src/jmh/java/datadog/trace/instrumentation/jdbc/SQLCommenterBenchmark.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ public void testToComment() {
dbService,
hostname,
dbName,
null,
env,
version,
traceParent);
Expand Down
21 changes: 21 additions & 0 deletions ...t/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/SQLCommenter.java
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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");
Expand Down Expand Up @@ -94,13 +99,21 @@ 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();
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 peerService = peerServiceObj != null ? peerServiceObj.toString() : null;

if (appendComment) {
sb.append(sql);
sb.append(SPACE);
Expand All @@ -113,6 +126,7 @@ public static String inject(
dbService,
hostname,
dbName,
peerService,
env,
version,
traceParent);
Expand All @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down
52 changes: 45 additions & 7 deletions dd-java-agent/instrumentation/jdbc/src/test/groovy/SQLCommenterTest.groovy
View file
Original file line number Diff line number Diff line change
@@ -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"() {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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'*/"
}
}

0 comments on commit 6ec8455

Please sign in to comment.