Skip to content

Commit 9f6ef81

Browse files
na-jina-ji
committed
Add option to JDBC instrumentation to always append DBM comment
DBMON-5791
1 parent c60b3c3 commit 9f6ef81

File tree

8 files changed

+66
-5
lines changed

8 files changed

+66
-5
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,8 @@ public static String onEnter(
126126
dbService =
127127
traceConfig(activeSpan()).getServiceMapping().getOrDefault(dbService, dbService);
128128
}
129-
boolean append = "sqlserver".equals(dbInfo.getType());
129+
130+
boolean append = "sqlserver".equals(dbInfo.getType()) || DECORATE.shouldAppendSqlComment();
130131
sql =
131132
SQLCommenter.inject(
132133
sql, dbService, dbInfo.getType(), dbInfo.getHost(), dbInfo.getDb(), null, append);

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,4 +423,8 @@ public boolean shouldInjectSQLComment() {
423423
return Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_FULL)
424424
|| Config.get().getDbmPropagationMode().equals(DBM_PROPAGATION_MODE_STATIC);
425425
}
426+
427+
public boolean shouldAppendSqlComment() {
428+
return Config.get().isDbmAlwaysAppendSqlComment();
429+
}
426430
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,14 +127,15 @@ public static AgentScope onEnter(
127127
final boolean injectTraceInComment = injectTraceContext && !isSqlServer && !isOracle;
128128

129129
// prepend mode will prepend the SQL comment to the raw sql query
130-
boolean appendComment = false;
130+
boolean appendComment = DECORATE.shouldAppendSqlComment();
131131

132132
// There is a bug in the SQL Server JDBC driver that prevents
133133
// the generated keys from being returned when the
134134
// SQL comment is prepended to the SQL query.
135135
// We only append in this case to avoid the comment from being truncated.
136136
// @see https://github.com/microsoft/mssql-jdbc/issues/2729
137137
if (isSqlServer
138+
&& !appendComment
138139
&& args.length == 2
139140
&& args[1] instanceof Integer
140141
&& (Integer) args[1] == Statement.RETURN_GENERATED_KEYS) {

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ class DBMInjectionForkedTest extends InstrumentationSpecification {
2121

2222
static query = "SELECT 1"
2323
static serviceInjection = "ddps='my_service_name',dddbs='remapped_testdb'"
24-
static fullInjection = serviceInjection + ",traceparent='00-00000000000000000000000000000004-0000000000000003-01'"
2524

2625
def "prepared stmt"() {
2726
setup:
@@ -36,6 +35,20 @@ class DBMInjectionForkedTest extends InstrumentationSpecification {
3635
assert statement.sql == "/*${serviceInjection}*/ ${query}"
3736
}
3837

38+
def "append comment on prepared stmt"() {
39+
setup:
40+
injectSysConfig(TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT, "true")
41+
def connection = new TestConnection(false)
42+
43+
when:
44+
def statement = connection.prepareStatement(query) as TestPreparedStatement
45+
statement.execute()
46+
47+
then:
48+
// even in full propagation mode, we cannot inject trace info in prepared statements
49+
assert statement.sql == "${query} /*${serviceInjection}*/"
50+
}
51+
3952
def "single query"() {
4053
setup:
4154
def connection = new TestConnection(false)
@@ -45,6 +58,19 @@ class DBMInjectionForkedTest extends InstrumentationSpecification {
4558
statement.executeQuery(query)
4659

4760
then:
48-
assert statement.sql == "/*${fullInjection}*/ ${query}"
61+
assert statement.sql == "/*${serviceInjection},traceparent='00-00000000000000000000000000000006-0000000000000005-01'*/ ${query}"
62+
}
63+
64+
def "append comment on single query"() {
65+
setup:
66+
injectSysConfig(TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT, "true")
67+
def connection = new TestConnection(false)
68+
69+
when:
70+
def statement = connection.createStatement() as TestStatement
71+
statement.executeQuery(query)
72+
73+
then:
74+
assert statement.sql == "${query} /*${serviceInjection},traceparent='00-00000000000000000000000000000008-0000000000000007-01'*/"
4975
}
5076
}

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ class SQLServerInjectionForkedTest extends InstrumentationSpecification {
3737
assert !statement.sql.contains("traceparent")
3838
}
3939

40-
def "SQL Server apend comment when getting generated keys"() {
40+
def "SQL Server append comment when getting generated keys"() {
4141
setup:
4242
def connection = new TestConnection(false)
4343
def metadata = new TestDatabaseMetaData()
@@ -51,4 +51,20 @@ class SQLServerInjectionForkedTest extends InstrumentationSpecification {
5151
then:
5252
assert statement.sql == "${query} /*${serviceInjection}*/"
5353
}
54+
55+
def "SQL Server append comment when configured to do so"() {
56+
setup:
57+
injectSysConfig(TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT, "true")
58+
def connection = new TestConnection(false)
59+
def metadata = new TestDatabaseMetaData()
60+
metadata.setURL("jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=testdb;")
61+
connection.setMetaData(metadata)
62+
63+
when:
64+
def statement = connection.createStatement() as TestStatement
65+
statement.executeUpdate(query)
66+
67+
then:
68+
assert statement.sql == "${query} /*${serviceInjection}*/"
69+
}
5470
}

‎dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ public final class ConfigDefaults {
7272
static final boolean DEFAULT_DB_CLIENT_HOST_SPLIT_BY_HOST = false;
7373
static final String DEFAULT_DB_DBM_PROPAGATION_MODE_MODE = "disabled";
7474
static final boolean DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS = false;
75+
static final boolean DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT = false;
7576
// Default value is set to 0, it disables the latency trace interceptor
7677
static final int DEFAULT_TRACE_KEEP_LATENCY_THRESHOLD_MS = 0;
7778
static final int DEFAULT_SCOPE_DEPTH_LIMIT = 100;

‎dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ public final class TraceInstrumentationConfig {
6969
public static final String DB_DBM_INJECT_SQL_BASEHASH = "dbm.inject.sql.basehash";
7070
public static final String DB_DBM_PROPAGATION_MODE_MODE = "dbm.propagation.mode";
7171
public static final String DB_DBM_TRACE_PREPARED_STATEMENTS = "dbm.trace_prepared_statements";
72+
public static final String DB_DBM_ALWAYS_APPEND_SQL_COMMENT = "dbm.always_append_sql_comment";
7273

7374
public static final String JDBC_CONNECTION_CLASS_NAME = "trace.jdbc.connection.class.name";
7475

‎internal-api/src/main/java/datadog/trace/api/Config.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_CLIENT_HOST_SPLIT_BY_INSTANCE_TYPE_SUFFIX;
5858
import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_PROPAGATION_MODE_MODE;
5959
import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS;
60+
import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT;
6061
import static datadog.trace.api.ConfigDefaults.DEFAULT_DEBUGGER_EXCEPTION_CAPTURE_INTERMEDIATE_SPANS_ENABLED;
6162
import static datadog.trace.api.ConfigDefaults.DEFAULT_DEBUGGER_EXCEPTION_CAPTURE_INTERVAL_SECONDS;
6263
import static datadog.trace.api.ConfigDefaults.DEFAULT_DEBUGGER_EXCEPTION_ENABLED;
@@ -511,6 +512,7 @@
511512
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_INJECT_SQL_BASEHASH;
512513
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_PROPAGATION_MODE_MODE;
513514
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_TRACE_PREPARED_STATEMENTS;
515+
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_DBM_ALWAYS_APPEND_SQL_COMMENT;
514516
import static datadog.trace.api.config.TraceInstrumentationConfig.ELASTICSEARCH_BODY_AND_PARAMS_ENABLED;
515517
import static datadog.trace.api.config.TraceInstrumentationConfig.ELASTICSEARCH_BODY_ENABLED;
516518
import static datadog.trace.api.config.TraceInstrumentationConfig.ELASTICSEARCH_PARAMS_ENABLED;
@@ -1076,6 +1078,7 @@ public static String getHostName() {
10761078
private final boolean dbmInjectSqlBaseHash;
10771079
private final String dbmPropagationMode;
10781080
private final boolean dbmTracePreparedStatements;
1081+
private final boolean dbmAlwaysAppendSqlComment;
10791082

10801083
private final boolean dynamicInstrumentationEnabled;
10811084
private final String dynamicInstrumentationSnapshotUrl;
@@ -1638,6 +1641,10 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
16381641
configProvider.getBoolean(
16391642
DB_DBM_TRACE_PREPARED_STATEMENTS, DEFAULT_DB_DBM_TRACE_PREPARED_STATEMENTS);
16401643

1644+
dbmAlwaysAppendSqlComment =
1645+
configProvider.getBoolean(
1646+
DB_DBM_ALWAYS_APPEND_SQL_COMMENT, DEFAULT_DB_DBM_ALWAYS_APPEND_SQL_COMMENT);
1647+
16411648
dbmInjectSqlBaseHash = configProvider.getBoolean(DB_DBM_INJECT_SQL_BASEHASH, false);
16421649

16431650
splitByTags = tryMakeImmutableSet(configProvider.getList(SPLIT_BY_TAGS));
@@ -5035,6 +5042,10 @@ public boolean isDbmTracePreparedStatements() {
50355042
return dbmTracePreparedStatements;
50365043
}
50375044

5045+
public boolean isDbmAlwaysAppendSqlComment() {
5046+
return dbmAlwaysAppendSqlComment;
5047+
}
5048+
50385049
public String getDbmPropagationMode() {
50395050
return dbmPropagationMode;
50405051
}

0 commit comments

Comments
 (0)