Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
35 commits
Select commit Hold shift + click to select a range
18cab70
Added support for DBM comment injection with MongoDB
Yoone Sep 22, 2025
4c5095b
Moved SharedDBCommenter from jdbc instrumentation to dd-trace-core
Yoone Sep 23, 2025
48c1358
Fixed SharedDBCommenterTest import
Yoone Sep 23, 2025
b3e777a
Replaced hardcoded length with constant
Yoone Sep 23, 2025
a26c412
Attempt at fixing existing tests
Yoone Sep 23, 2025
5df1948
Updated INJECT_COMMENT logic
Yoone Sep 23, 2025
486e238
Fixed containsTraceComment for test cases
Yoone Sep 24, 2025
29aef39
Switched MongoDBMInjectionTest to MongoBaseTest
Yoone Sep 24, 2025
7515a3f
Revamped Mongo comment tests
Yoone Sep 24, 2025
2533add
Attempt to fix SQL injection tests
Yoone Oct 6, 2025
dfa23a6
Revert changes to gradle.lockfile
Yoone Oct 6, 2025
ee8f5a0
PR comments: build traceparent incl. sampling flag, exception logging
Yoone Oct 14, 2025
7da62c3
Add missing helperClassNames for mongo drivers 3.1 and 3.4
na-ji Oct 15, 2025
e8e0a3d
Using unit testing for trace parent changes
Yoone Oct 15, 2025
a0d854f
Prevented cloning immutable Mongo objects
Yoone Oct 15, 2025
92a6054
Fixed smoke test (revert bug)
Yoone Oct 15, 2025
3f8198c
add new mongo instrumentation to inject APM/DBM comment
na-ji Oct 29, 2025
20eeee0
add support for mongo driver version 4.3
na-ji Nov 13, 2025
5c7422c
Merge remote-tracking branch 'origin/master' into yoann.bentz/mongodb…
na-ji Nov 18, 2025
4f6ca55
Merge remote-tracking branch 'origin/master' into yoann.bentz/mongodb…
na-ji Dec 3, 2025
0a2897b
rename MongoCommentInjector's getComment to buildComment
na-ji Dec 4, 2025
5476cd6
only create a mutable copy of the BsonDocument when needed
na-ji Dec 5, 2025
d2a91cc
move shared sql commenter to bootstrap package
na-ji Dec 5, 2025
fa7d331
document why we check for existing span on command listener
na-ji Dec 8, 2025
35ee37f
reformat trace parent util
na-ji Dec 8, 2025
9da2307
fix server connection instrumentation advice not working on write que…
na-ji Dec 10, 2025
81c02a7
Merge remote-tracking branch 'origin/master' into yoann.bentz/mongodb…
na-ji Dec 15, 2025
8a36080
improve check for overload instrumentation
na-ji Dec 15, 2025
b3663d5
Merge remote-tracking branch 'origin/master' into yoann.bentz/mongodb…
na-ji Dec 15, 2025
b538f1c
Merge remote-tracking branch 'origin/master' into yoann.bentz/mongodb…
na-ji Dec 16, 2025
b6f781d
add missing helper class for mongo 3.1 instrumentation
na-ji Dec 16, 2025
5f3fdd5
move mongo < 4.0 to standalone package driver-3.6
na-ji Dec 16, 2025
7cd4d47
add mock method to force the DefaultServerConnection40Instrumentation…
na-ji Dec 16, 2025
047a79a
Move DefaultServerInstrumentation for mongo driver 4.0 to driver 3.8
na-ji Dec 17, 2025
5ffb54e
Merge remote-tracking branch 'origin/master' into yoann.bentz/mongodb…
na-ji Dec 17, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .github/CODEOWNERS
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,9 @@ dd-java-agent/agent-llmobs/ @DataDog/ml-observability
dd-trace-core/src/main/java/datadog/trace/llmobs/ @DataDog/ml-observability
dd-trace-core/src/test/groovy/datadog/trace/llmobs/ @DataDog/ml-observability

# @DataDog/database-monitoring
datadog/trace/bootstrap/instrumentation/dbm @DataDog/database-monitoring @DataDog/apm-idm-java

# @DataDog/rum
/internal-api/src/main/java/datadog/trace/api/rum/ @DataDog/rum
/internal-api/src/test/groovy/datadog/trace/api/rum/ @DataDog/rum
Expand Down
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,10 @@ out/
##########
.cursor

# Vim #
#######
*.sw[nop]

# Others #
##########
/logs/*
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
package datadog.trace.bootstrap.instrumentation.dbm;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;

import datadog.trace.api.BaseHash;
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;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Shared database comment builder for generating trace context comments for SQL DBs and MongoDB */
public class SharedDBCommenter {
private static final Logger log = LoggerFactory.getLogger(SharedDBCommenter.class);
private static final String UTF8 = StandardCharsets.UTF_8.toString();

private static final char EQUALS = '=';
private static final char COMMA = ',';
private static final char QUOTE = '\'';

// Injected fields. When adding a new one, be sure to update this and the methods below.
private static final String PARENT_SERVICE = encode("ddps");
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");
private static final String DD_SERVICE_HASH = encode("ddsh");

// Used by SQLCommenter and MongoCommentInjector to avoid duplicate comment injection
public static boolean containsTraceComment(String commentContent) {
return commentContent.contains(PARENT_SERVICE + "=")
|| commentContent.contains(DATABASE_SERVICE + "=")
|| commentContent.contains(DD_HOSTNAME + "=")
|| commentContent.contains(DD_DB_NAME + "=")
|| commentContent.contains(DD_PEER_SERVICE + "=")
|| commentContent.contains(DD_ENV + "=")
|| commentContent.contains(DD_VERSION + "=")
|| commentContent.contains(TRACEPARENT + "=")
|| commentContent.contains(DD_SERVICE_HASH + "=");
Comment on lines +37 to +45
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know you only moved that code, but I wonder after how many "contains" does a precompiled regex becomes more performant.

}

// Build database comment content without comment delimiters such as /* */
public static String buildComment(
String dbService, String dbType, String hostname, String dbName, String traceParent) {

Config config = Config.get();
StringBuilder sb = new StringBuilder();

int initSize = 0; // No initial content for pure comment
append(sb, PARENT_SERVICE, config.getServiceName(), initSize);
append(sb, DATABASE_SERVICE, dbService, initSize);
append(sb, DD_HOSTNAME, hostname, initSize);
append(sb, DD_DB_NAME, dbName, initSize);
append(sb, DD_PEER_SERVICE, getPeerService(), initSize);
append(sb, DD_ENV, config.getEnv(), initSize);
append(sb, DD_VERSION, config.getVersion(), initSize);
append(sb, TRACEPARENT, traceParent, initSize);

if (config.isDbmInjectSqlBaseHash() && config.isExperimentalPropagateProcessTagsEnabled()) {
append(sb, DD_SERVICE_HASH, BaseHash.getBaseHashStr(), initSize);
}

return sb.length() > 0 ? sb.toString() : null;
}

private static String getPeerService() {
AgentSpan span = activeSpan();
Object peerService = null;
if (span != null) {
peerService = span.getTag(Tags.PEER_SERVICE);
}
return peerService != null ? peerService.toString() : null;
}

private static String encode(String val) {
try {
return URLEncoder.encode(val, UTF8);
} catch (UnsupportedEncodingException exe) {
if (log.isDebugEnabled()) {
log.debug("exception thrown while encoding comment key {}", val, exe);
}
}
return val;
}

private static void append(StringBuilder sb, String key, String value, int initSize) {
if (null == value || value.isEmpty()) {
return;
}
String encodedValue;
try {
encodedValue = URLEncoder.encode(value, UTF8);
} catch (UnsupportedEncodingException e) {
encodedValue = value;
}
if (sb.length() > initSize) {
sb.append(COMMA);
}
sb.append(key).append(EQUALS).append(QUOTE).append(encodedValue).append(QUOTE);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT;
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.logQueryInfoInjection;
import static net.bytebuddy.matcher.ElementMatchers.returns;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
Expand Down Expand Up @@ -77,7 +78,9 @@ public DBMCompatibleConnectionInstrumentation() {
@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".JDBCDecorator", packageName + ".SQLCommenter",
packageName + ".JDBCDecorator",
packageName + ".SQLCommenter",
"datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter",
};
}

Expand Down Expand Up @@ -110,8 +113,7 @@ public static class ConnectionAdvice {
public static String onEnter(
@Advice.This Connection connection,
@Advice.Argument(value = 0, readOnly = false) String sql) {
// Using INJECT_COMMENT fails to update when a test calls injectSysConfig
if (!DECORATE.shouldInjectSQLComment()) {
if (!INJECT_COMMENT) {
return sql;
}
if (CallDepthThreadLocalMap.incrementCallDepth(Connection.class) > 0) {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
package datadog.trace.instrumentation.jdbc;

import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_FULL;
import static datadog.trace.api.Config.DBM_PROPAGATION_MODE_STATIC;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.INSTRUMENTATION_TIME_MS;
Expand Down Expand Up @@ -48,8 +50,6 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
UTF8BytesString.create("java-jdbc-prepared_statement");
private static final String DEFAULT_SERVICE_NAME =
SpanNaming.instance().namingSchema().database().service("jdbc");
public static final String DBM_PROPAGATION_MODE_STATIC = "service";
public static final String DBM_PROPAGATION_MODE_FULL = "full";

public static final String DD_INSTRUMENTATION_PREFIX = "_DD_";

Expand Down Expand Up @@ -417,11 +417,6 @@ public boolean shouldInjectTraceContext(DBInfo dbInfo) {
return INJECT_TRACE_CONTEXT;
}

public boolean shouldInjectSQLComment() {
return Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_FULL)
|| Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_STATIC);
}

private void logInjectionErrorOnce(String vessel, Throwable t) {
if (!loggedInjectionError) {
loggedInjectionError = true;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,47 +1,19 @@
package datadog.trace.instrumentation.jdbc;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;

import datadog.trace.api.BaseHash;
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;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter;

public class SQLCommenter {
private static final Logger log = LoggerFactory.getLogger(SQLCommenter.class);
private static final String UTF8 = StandardCharsets.UTF_8.toString();

private static final char EQUALS = '=';
private static final char COMMA = ',';
private static final char QUOTE = '\'';
// SQL-specific constants, rest defined in SharedDBCommenter
private static final char SPACE = ' ';
private static final String OPEN_COMMENT = "/*";
private static final int OPEN_COMMENT_LEN = OPEN_COMMENT.length();
private static final String CLOSE_COMMENT = "*/";

// Injected fields. When adding a new one, be sure to update this and the methods below.
private static final int NUMBER_OF_FIELDS = 9;
private static final String PARENT_SERVICE = encode("ddps");
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");
private static final String DD_SERVICE_HASH = encode("ddsh");

private static final int KEY_AND_SEPARATORS_ESTIMATED_SIZE = 10;
private static final int VALUE_ESTIMATED_SIZE = 10;
private static final int TRACE_PARENT_EXTRA_ESTIMATED_SIZE = 50;
private static final int INJECTED_COMMENT_ESTIMATED_SIZE =
NUMBER_OF_FIELDS * (KEY_AND_SEPARATORS_ESTIMATED_SIZE + VALUE_ESTIMATED_SIZE)
+ TRACE_PARENT_EXTRA_ESTIMATED_SIZE;
// Size estimation for StringBuilder pre-allocation
private static final int SPACE_CHARS = 2; // Leading and trailing spaces
private static final int COMMENT_DELIMITERS = 4; // "/*" + "*/"
private static final int BUFFER_EXTRA = 4;
private static final int SQL_COMMENT_OVERHEAD = SPACE_CHARS + COMMENT_DELIMITERS + BUFFER_EXTRA;

protected static String getFirstWord(String sql) {
int beginIndex = 0;
Expand Down Expand Up @@ -99,9 +71,16 @@ public static String inject(
return sql;
}

Config config = Config.get();
String commentContent =
SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent);

if (commentContent == null) {
return sql;
}

StringBuilder sb = new StringBuilder(sql.length() + INJECTED_COMMENT_ESTIMATED_SIZE);
// SQL-specific wrapping with /* */
StringBuilder sb =
new StringBuilder(sql.length() + commentContent.length() + SQL_COMMENT_OVERHEAD);
int closingSemicolonIndex = indexOfClosingSemicolon(sql);
if (appendComment) {
if (closingSemicolonIndex > -1) {
Expand All @@ -113,22 +92,7 @@ public static String inject(
}

sb.append(OPEN_COMMENT);
int initSize = sb.length();
append(sb, PARENT_SERVICE, config.getServiceName(), initSize);
append(sb, DATABASE_SERVICE, dbService, initSize);
append(sb, DD_HOSTNAME, hostname, initSize);
append(sb, DD_DB_NAME, dbName, initSize);
append(sb, DD_PEER_SERVICE, getPeerService(), initSize);
append(sb, DD_ENV, config.getEnv(), initSize);
append(sb, DD_VERSION, config.getVersion(), initSize);
append(sb, TRACEPARENT, traceParent, initSize);
if (config.isDbmInjectSqlBaseHash() && config.isExperimentalPropagateProcessTagsEnabled()) {
append(sb, DD_SERVICE_HASH, BaseHash.getBaseHashStr(), initSize);
}
if (initSize == sb.length()) {
// no comment was added
return sql;
}
sb.append(commentContent);
sb.append(CLOSE_COMMENT);
if (!appendComment) {
sb.append(SPACE);
Expand All @@ -142,72 +106,30 @@ public static String inject(
return sb.toString();
}

private static String getPeerService() {
AgentSpan span = activeSpan();
Object peerService = null;
if (span != null) {
peerService = span.getTag(Tags.PEER_SERVICE);
}
return peerService != null ? peerService.toString() : null;
}

private static boolean hasDDComment(String sql, boolean appendComment) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this is called for all SQL requests, I wonder if we could make it more lightweight

if ((!sql.endsWith(CLOSE_COMMENT) && appendComment)
|| ((!sql.startsWith(OPEN_COMMENT)) && !appendComment)) {
return false;
}
int startIdx = OPEN_COMMENT_LEN;
if (appendComment) {
startIdx += sql.lastIndexOf(OPEN_COMMENT);
}
return hasMatchingSubstring(sql, startIdx, PARENT_SERVICE)
|| hasMatchingSubstring(sql, startIdx, DATABASE_SERVICE)
|| hasMatchingSubstring(sql, startIdx, DD_HOSTNAME)
|| hasMatchingSubstring(sql, startIdx, DD_DB_NAME)
|| hasMatchingSubstring(sql, startIdx, DD_PEER_SERVICE)
|| hasMatchingSubstring(sql, startIdx, DD_ENV)
|| hasMatchingSubstring(sql, startIdx, DD_VERSION)
|| hasMatchingSubstring(sql, startIdx, TRACEPARENT)
|| hasMatchingSubstring(sql, startIdx, DD_SERVICE_HASH);
}

private static boolean hasMatchingSubstring(String sql, int startIndex, String substring) {
boolean tooLong = startIndex + substring.length() >= sql.length();
if (tooLong || !(sql.charAt(startIndex + substring.length()) == EQUALS)) {
return false;
}
return sql.startsWith(substring, startIndex);
String commentContent = extractCommentContent(sql, appendComment);
return SharedDBCommenter.containsTraceComment(commentContent);
}

private static String encode(String val) {
try {
return URLEncoder.encode(val, UTF8);
} catch (UnsupportedEncodingException exe) {
if (log.isDebugEnabled()) {
log.debug("exception thrown while encoding sql comment key %s", exe);
}
}
return val;
}

private static void append(StringBuilder sb, String key, String value, int initSize) {
if (null == value || value.isEmpty()) {
return;
}
String encodedValue;
try {
encodedValue = URLEncoder.encode(value, UTF8);
} catch (UnsupportedEncodingException e) {
if (log.isDebugEnabled()) {
log.debug("exception thrown while encoding sql comment %s", e);
}
return;
private static String extractCommentContent(String sql, boolean appendComment) {
int startIdx;
int endIdx;
if (appendComment) {
startIdx = sql.lastIndexOf(OPEN_COMMENT);
endIdx = sql.lastIndexOf(CLOSE_COMMENT);
} else {
startIdx = sql.indexOf(OPEN_COMMENT);
endIdx = sql.indexOf(CLOSE_COMMENT);
}

if (sb.length() > initSize) {
sb.append(COMMA);
if (startIdx != -1 && endIdx != -1 && endIdx > startIdx) {
return sql.substring(startIdx + OPEN_COMMENT_LEN, endIdx);
}
sb.append(key).append(EQUALS).append(QUOTE).append(encodedValue).append(QUOTE);
return "";
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,11 @@ public Map<String, String> contextStore() {

@Override
public String[] helperClassNames() {
return new String[] {packageName + ".JDBCDecorator", packageName + ".SQLCommenter"};
return new String[] {
packageName + ".JDBCDecorator",
packageName + ".SQLCommenter",
"datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter",
};
}

@Override
Expand Down
4 changes: 4 additions & 0 deletions dd-java-agent/instrumentation/mongo/common/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -2,5 +2,9 @@ apply from: "$rootDir/gradle/java.gradle"

dependencies {
compileOnly group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'

implementation project(':dd-trace-core')

testImplementation project(':dd-java-agent:instrumentation:mongo').sourceSets.test.output
testImplementation group: 'org.mongodb', name: 'mongo-java-driver', version: '3.1.0'
}
Loading