Skip to content

Commit abee429

Browse files
nenadnoveljicnenadnoveljic
committed
Add DBM_TRACE_INJECTED tag to SQL Server prepared statements
1 parent 104a441 commit abee429

3 files changed

Lines changed: 13 additions & 5 deletions

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DBM_TRACE_INJECTED;
67
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DATABASE_QUERY;
78
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.DECORATE;
89
import static datadog.trace.instrumentation.jdbc.JDBCDecorator.INJECT_COMMENT;
@@ -84,6 +85,7 @@ public static AgentScope onEnter(@Advice.This final Statement statement) {
8485
final long spanID = DECORATE.setContextInfo(connection, dbInfo);
8586
// we then force that pre-determined span ID for the span covering the actual query
8687
span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
88+
span.setTag(DBM_TRACE_INJECTED, true);
8789
} else {
8890
span = startSpan(DATABASE_QUERY);
8991
}

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -93,22 +93,21 @@ public static AgentScope onEnter(
9393
final long spanID = DECORATE.setContextInfo(connection, dbInfo);
9494
// we then force that pre-determined span ID for the span covering the actual query
9595
span = AgentTracer.get().buildSpan(DATABASE_QUERY).withSpanId(spanID).start();
96+
span.setTag(DBM_TRACE_INJECTED, true);
9697
} else {
9798
span = startSpan(DATABASE_QUERY);
9899
}
99100

100101
DECORATE.afterStart(span);
101102
DECORATE.onConnection(span, dbInfo);
102103
final String copy = sql;
103-
if (span != null && INJECT_COMMENT) {
104+
if (span != null && INJECT_COMMENT && !isSqlServer) {
104105
String traceParent = null;
105106

106107
if (injectTraceContext) {
107108
Integer priority = span.forceSamplingDecision();
108109
if (priority != null) {
109-
if (!isSqlServer) {
110-
traceParent = DECORATE.traceParent(span, priority);
111-
}
110+
traceParent = DECORATE.traceParent(span, priority);
112111
// set the dbm trace injected tag on the span
113112
span.setTag(DBM_TRACE_INJECTED, true);
114113
}

‎dd-java-agent/instrumentation/jdbc/src/test/groovy/RemoteJDBCInstrumentationTest.groovy‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -346,7 +346,8 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
346346
then:
347347
resultSet.next()
348348
resultSet.getInt(1) == 3
349-
if (driver == POSTGRESQL || driver == MYSQL || !dbmTraceInjected()) {
349+
def addDbmTag = dbmTraceInjected()
350+
if (driver == POSTGRESQL || driver == MYSQL || !addDbmTag) {
350351
assertTraces(1) {
351352
trace(2) {
352353
basicSpan(it, "parent")
@@ -402,6 +403,9 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
402403
if (usingHikari) {
403404
"$Tags.DB_POOL_NAME" String
404405
}
406+
if (addDbmTag) {
407+
"$InstrumentationTags.DBM_TRACE_INJECTED" true
408+
}
405409
peerServiceFrom(Tags.DB_INSTANCE)
406410
defaultTags()
407411
}
@@ -526,6 +530,7 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
526530
if (conPoolType == "hikari") {
527531
"$Tags.DB_POOL_NAME" String
528532
}
533+
"$InstrumentationTags.DBM_TRACE_INJECTED" true
529534
peerServiceFrom(Tags.DB_INSTANCE)
530535
defaultTags()
531536
}
@@ -645,6 +650,7 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
645650
if (conPoolType == "hikari") {
646651
"$Tags.DB_POOL_NAME" String
647652
}
653+
"$InstrumentationTags.DBM_TRACE_INJECTED" true
648654
defaultTags()
649655
}
650656
}
@@ -774,6 +780,7 @@ abstract class RemoteJDBCInstrumentationTest extends VersionedNamingTestBase {
774780
if (conPoolType == "hikari") {
775781
"$Tags.DB_POOL_NAME" String
776782
}
783+
"$InstrumentationTags.DBM_TRACE_INJECTED" true
777784
peerServiceFrom(Tags.DB_INSTANCE)
778785
defaultTags()
779786
}

0 commit comments

Comments
 (0)