Skip to content

Commit 95550ac

Browse files
author
Andrew Kent
committed
Save connection info after connection construction.
1 parent 28bc1e7 commit 95550ac

File tree

5 files changed

+76
-101
lines changed

5 files changed

+76
-101
lines changed

‎dd-java-agent-ittests/src/test/groovy/com/datadoghq/agent/integration/jdbc/JDBCInstrumentationTest.groovy‎

Lines changed: 33 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package com.datadoghq.agent.integration.jdbc
22

33
import com.datadoghq.trace.DDTracer
44
import com.datadoghq.trace.writer.ListWriter
5+
import dd.test.TestUtils
56
import io.opentracing.util.GlobalTracer
67
import org.apache.derby.jdbc.EmbeddedDriver
78
import org.h2.Driver
@@ -18,8 +19,8 @@ import java.sql.Statement
1819

1920
class JDBCInstrumentationTest extends Specification {
2021

21-
ListWriter writer = new ListWriter()
22-
DDTracer tracer = new DDTracer(writer)
22+
final ListWriter writer = new ListWriter()
23+
final DDTracer tracer = new DDTracer(writer)
2324

2425
@Shared
2526
private Map<String, Connection> connections
@@ -43,16 +44,8 @@ class JDBCInstrumentationTest extends Specification {
4344
}
4445

4546
def setup() {
46-
try {
47-
GlobalTracer.register(tracer)
48-
} catch (final Exception e) {
49-
// Force it anyway using reflection
50-
final Field field = GlobalTracer.getDeclaredField("tracer")
51-
field.setAccessible(true)
52-
field.set(null, tracer)
53-
}
47+
TestUtils.registerOrReplaceGlobalTracer(tracer)
5448
writer.start()
55-
assert GlobalTracer.isRegistered()
5649
}
5750

5851
@Unroll
@@ -80,6 +73,7 @@ class JDBCInstrumentationTest extends Specification {
8073

8174
def tags = span.context().tags
8275
tags["db.type"] == driver
76+
tags["db.user"] == username
8377
tags["span.kind"] == "client"
8478
tags["component"] == "java-jdbc-statement"
8579

@@ -88,16 +82,16 @@ class JDBCInstrumentationTest extends Specification {
8882

8983
tags["thread.name"] != null
9084
tags["thread.id"] != null
91-
tags.size() == 7
85+
tags.size() == username == null ? 7 : 8
9286

9387
cleanup:
9488
statement.close()
9589

9690
where:
97-
driver | connection | query
98-
"h2" | connections.get("h2") | "SELECT 3"
99-
"derby" | connections.get("derby") | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
100-
"hsqldb" | connections.get("hsqldb") | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
91+
driver | connection | username | query
92+
"h2" | connections.get("h2") | null | "SELECT 3"
93+
"derby" | connections.get("derby") | "APP" | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
94+
"hsqldb" | connections.get("hsqldb") | "SA" | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
10195
}
10296

10397
@Unroll
@@ -126,6 +120,7 @@ class JDBCInstrumentationTest extends Specification {
126120

127121
def tags = span.context().tags
128122
tags["db.type"] == driver
123+
tags["db.user"] == username
129124
tags["span.kind"] == "client"
130125
tags["component"] == "java-jdbc-prepared_statement"
131126

@@ -134,16 +129,16 @@ class JDBCInstrumentationTest extends Specification {
134129

135130
tags["thread.name"] != null
136131
tags["thread.id"] != null
137-
tags.size() == 7
132+
tags.size() == username == null ? 7 : 8
138133

139134
cleanup:
140135
statement.close()
141136

142137
where:
143-
driver | connection | query
144-
"h2" | connections.get("h2") | "SELECT 3"
145-
"derby" | connections.get("derby") | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
146-
"hsqldb" | connections.get("hsqldb") | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
138+
driver | connection | username | query
139+
"h2" | connections.get("h2") | null | "SELECT 3"
140+
"derby" | connections.get("derby") | "APP" | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
141+
"hsqldb" | connections.get("hsqldb") | "SA" | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
147142
}
148143

149144
@Unroll
@@ -171,6 +166,7 @@ class JDBCInstrumentationTest extends Specification {
171166

172167
def tags = span.context().tags
173168
tags["db.type"] == driver
169+
tags["db.user"] == username
174170
tags["span.kind"] == "client"
175171
tags["component"] == "java-jdbc-prepared_statement"
176172

@@ -179,16 +175,16 @@ class JDBCInstrumentationTest extends Specification {
179175

180176
tags["thread.name"] != null
181177
tags["thread.id"] != null
182-
tags.size() == 7
178+
tags.size() == username == null ? 7 : 8
183179

184180
cleanup:
185181
statement.close()
186182

187183
where:
188-
driver | connection | query
189-
"h2" | connections.get("h2") | "SELECT 3"
190-
"derby" | connections.get("derby") | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
191-
"hsqldb" | connections.get("hsqldb") | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
184+
driver | connection | username | query
185+
"h2" | connections.get("h2") | null | "SELECT 3"
186+
"derby" | connections.get("derby") | "APP" | "SELECT 3 FROM SYSIBM.SYSDUMMY1"
187+
"hsqldb" | connections.get("hsqldb") | "SA" | "SELECT 3 FROM INFORMATION_SCHEMA.SYSTEM_USERS"
192188
}
193189

194190
@Unroll
@@ -217,6 +213,7 @@ class JDBCInstrumentationTest extends Specification {
217213

218214
def tags = span.context().tags
219215
tags["db.type"] == driver
216+
tags["db.user"] == username
220217
tags["span.kind"] == "client"
221218
tags["component"] == "java-jdbc-statement"
222219

@@ -225,16 +222,16 @@ class JDBCInstrumentationTest extends Specification {
225222

226223
tags["thread.name"] != null
227224
tags["thread.id"] != null
228-
tags.size() == 7
225+
tags.size() == username == null ? 7 : 8
229226

230227
cleanup:
231228
statement.close()
232229

233230
where:
234-
driver | connection | query
235-
"h2" | connections.get("h2") | "CREATE TABLE S_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
236-
"derby" | connections.get("derby") | "CREATE TABLE S_DERBY (id INTEGER not NULL, PRIMARY KEY ( id ))"
237-
"hsqldb" | connections.get("hsqldb") | "CREATE TABLE PUBLIC.S_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
231+
driver | connection | username | query
232+
"h2" | connections.get("h2") | null | "CREATE TABLE S_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
233+
"derby" | connections.get("derby") | "APP" | "CREATE TABLE S_DERBY (id INTEGER not NULL, PRIMARY KEY ( id ))"
234+
"hsqldb" | connections.get("hsqldb") | "SA" | "CREATE TABLE PUBLIC.S_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
238235
}
239236

240237
@Unroll
@@ -261,6 +258,7 @@ class JDBCInstrumentationTest extends Specification {
261258

262259
def tags = span.context().tags
263260
tags["db.type"] == driver
261+
tags["db.user"] == username
264262
tags["span.kind"] == "client"
265263
tags["component"] == "java-jdbc-prepared_statement"
266264

@@ -269,15 +267,15 @@ class JDBCInstrumentationTest extends Specification {
269267

270268
tags["thread.name"] != null
271269
tags["thread.id"] != null
272-
tags.size() == 7
270+
tags.size() == username == null ? 7 : 8
273271

274272
cleanup:
275273
statement.close()
276274

277275
where:
278-
driver | connection | query
279-
"h2" | connections.get("h2") | "CREATE TABLE PS_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
276+
driver | connection | username | query
277+
"h2" | connections.get("h2") | null | "CREATE TABLE PS_H2 (id INTEGER not NULL, PRIMARY KEY ( id ))"
280278
// Derby calls executeLargeUpdate from executeUpdate thus generating a nested span breaking this test.
281-
"hsqldb" | connections.get("hsqldb") | "CREATE TABLE PUBLIC.PS_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
279+
"hsqldb" | connections.get("hsqldb") | "SA" | "CREATE TABLE PUBLIC.PS_HSQLDB (id INTEGER not NULL, PRIMARY KEY ( id ))"
282280
}
283281
}

‎dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/ConnectionInstrumentation.java‎

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.datadoghq.agent.instrumentation.jdbc;
22

33
import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
4+
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
45
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
56
import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
67
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -15,11 +16,13 @@
1516
import java.sql.PreparedStatement;
1617
import java.util.Map;
1718
import java.util.WeakHashMap;
19+
import lombok.Data;
1820
import net.bytebuddy.agent.builder.AgentBuilder;
1921
import net.bytebuddy.asm.Advice;
2022

2123
@AutoService(Instrumenter.class)
2224
public final class ConnectionInstrumentation implements Instrumenter {
25+
public static final Map<Connection, DBInfo> connectionInfo = new WeakHashMap<>();
2326
public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
2427

2528
@Override
@@ -32,15 +35,47 @@ public AgentBuilder instrument(final AgentBuilder agentBuilder) {
3235
nameStartsWith("prepare")
3336
.and(takesArgument(0, String.class))
3437
.and(returns(PreparedStatement.class)),
35-
ConnectionAdvice.class.getName()))
38+
ConnectionPrepareAdvice.class.getName()))
39+
.transform(
40+
DDAdvice.create().advice(isConstructor(), ConnectionConstructorAdvice.class.getName()))
3641
.asDecorator();
3742
}
3843

39-
public static class ConnectionAdvice {
44+
public static class ConnectionPrepareAdvice {
4045
@Advice.OnMethodExit(suppress = Throwable.class)
4146
public static void addDBInfo(
4247
@Advice.Argument(0) final String sql, @Advice.Return final PreparedStatement statement) {
4348
preparedStatements.put(statement, sql);
4449
}
4550
}
51+
52+
public static class ConnectionConstructorAdvice {
53+
@Advice.OnMethodExit(suppress = Throwable.class)
54+
public static void addDBInfo(@Advice.This final Connection connection) {
55+
try {
56+
final String url = connection.getMetaData().getURL();
57+
if (url != null) {
58+
// Remove end of url to prevent passwords from leaking:
59+
final String sanitizedURL = url.replaceAll("[?;].*", "");
60+
final String type = url.split(":")[1];
61+
String user = connection.getMetaData().getUserName();
62+
if (user != null && user.trim().equals("")) {
63+
user = null;
64+
}
65+
connectionInfo.put(connection, new DBInfo(sanitizedURL, type, user));
66+
}
67+
} catch (Throwable t) {
68+
// object may not be fully initialized.
69+
// calling constructor will populate map
70+
}
71+
}
72+
}
73+
74+
@Data
75+
public static class DBInfo {
76+
public static DBInfo UNKNOWN = new DBInfo("null", "unknown", null);
77+
private final String url;
78+
private final String type;
79+
private final String user;
80+
}
4681
}

‎dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/DriverInstrumentation.java‎

Lines changed: 0 additions & 60 deletions
This file was deleted.

‎dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/PreparedStatementInstrumentation.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,10 @@ public static ActiveSpan startSpan(@Advice.This final PreparedStatement statemen
5151
return NoopActiveSpanSource.NoopActiveSpan.INSTANCE;
5252
}
5353

54-
DriverInstrumentation.DBInfo dbInfo = DriverInstrumentation.connectionInfo.get(connection);
54+
ConnectionInstrumentation.DBInfo dbInfo =
55+
ConnectionInstrumentation.connectionInfo.get(connection);
5556
if (dbInfo == null) {
56-
dbInfo = DriverInstrumentation.DBInfo.UNKNOWN;
57+
dbInfo = ConnectionInstrumentation.DBInfo.UNKNOWN;
5758
}
5859

5960
final ActiveSpan span =

‎dd-java-agent/src/main/java/com/datadoghq/agent/instrumentation/jdbc/StatementInstrumentation.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,10 @@ public static ActiveSpan startSpan(
5050
return NoopActiveSpanSource.NoopActiveSpan.INSTANCE;
5151
}
5252

53-
DriverInstrumentation.DBInfo dbInfo = DriverInstrumentation.connectionInfo.get(connection);
53+
ConnectionInstrumentation.DBInfo dbInfo =
54+
ConnectionInstrumentation.connectionInfo.get(connection);
5455
if (dbInfo == null) {
55-
dbInfo = DriverInstrumentation.DBInfo.UNKNOWN;
56+
dbInfo = ConnectionInstrumentation.DBInfo.UNKNOWN;
5657
}
5758

5859
final ActiveSpan span =

0 commit comments

Comments
 (0)