Skip to content

Commit 9f9842b

Browse files
reftelreftel
committed
Add JDBI 3 plugin
1 parent 24e5d41 commit 9f9842b

File tree

10 files changed

+519
-2
lines changed

10 files changed

+519
-2
lines changed

‎brave-bom/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,11 @@
162162
<artifactId>brave-instrumentation-jersey-server-jakarta</artifactId>
163163
<version>${project.version}</version>
164164
</dependency>
165+
<dependency>
166+
<groupId>${project.groupId}</groupId>
167+
<artifactId>brave-instrumentation-jdbi3</artifactId>
168+
<version>${project.version}</version>
169+
</dependency>
165170
<dependency>
166171
<groupId>${project.groupId}</groupId>
167172
<artifactId>brave-instrumentation-jms</artifactId>

‎instrumentation/jdbi3/README.md‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# brave-instrumentation-jdbi3
2+
3+
This includes a JDBI 3 plugin that will report to Zipkin how long each
4+
statement takes, along with relevant tags like the query.
5+
6+
To use it, call the `installPlugin` on the `Jdbi` instance you want to
7+
instrument, or add the statement context listener manually like so:
8+
```
9+
jdbi.getConfig(SqlStatements.class)
10+
.addContextListener(new BraveStatementContextListener(tracing));
11+
```
12+
13+
The remote service name of the span is set to the hostname and port number of
14+
the database server, if available, and the URL scheme if not. If the database
15+
URL format allows it, you can add the `zipkinServiceName` query parameter to
16+
override the remote service name.
17+
18+
Bind variable values are not included in the traces, only the SQL statement
19+
with placeholders.

‎instrumentation/jdbi3/pom.xml‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
<?xml version="1.0"?>
2+
<!--
3+
4+
Copyright The OpenZipkin Authors
5+
SPDX-License-Identifier: Apache-2.0
6+
7+
-->
8+
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
9+
<parent>
10+
<groupId>io.zipkin.brave</groupId>
11+
<artifactId>brave-instrumentation-parent</artifactId>
12+
<version>6.3.0-SNAPSHOT</version>
13+
</parent>
14+
<modelVersion>4.0.0</modelVersion>
15+
16+
<artifactId>brave-instrumentation-jdbi3</artifactId>
17+
<name>Brave Instrumentation: JDBI3</name>
18+
19+
<properties>
20+
<!-- Matches Export-Package in bnd.bnd -->
21+
<module.name>brave.jdbi3</module.name>
22+
23+
<main.basedir>${project.basedir}/../..</main.basedir>
24+
25+
<jdbi3.version>3.49.1</jdbi3.version>
26+
</properties>
27+
28+
<dependencies>
29+
<dependency>
30+
<groupId>org.jdbi</groupId>
31+
<artifactId>jdbi3-core</artifactId>
32+
<version>${jdbi3.version}</version>
33+
</dependency>
34+
<dependency>
35+
<groupId>${project.groupId}</groupId>
36+
<artifactId>brave-tests</artifactId>
37+
<version>${project.version}</version>
38+
<scope>test</scope>
39+
</dependency>
40+
<dependency>
41+
<groupId>com.mysql</groupId>
42+
<artifactId>mysql-connector-j</artifactId>
43+
<version>${mysql-connector-j8.version}</version>
44+
<scope>test</scope>
45+
</dependency>
46+
</dependencies>
47+
</project>

‎instrumentation/jdbi3/src/main/java/brave/jdbi3/BraveStatementContextListener.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
* Copyright The OpenZipkin Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package brave.jdbi3;
6+
7+
import java.net.URI;
8+
import java.net.URISyntaxException;
9+
import java.sql.Connection;
10+
import java.sql.SQLException;
11+
import java.time.Instant;
12+
import java.util.Map;
13+
import java.util.concurrent.ConcurrentHashMap;
14+
import java.util.regex.Matcher;
15+
import java.util.regex.Pattern;
16+
17+
import org.jdbi.v3.core.statement.SqlStatements;
18+
import org.jdbi.v3.core.statement.StatementContext;
19+
import org.jdbi.v3.core.statement.StatementContextListener;
20+
21+
import brave.Span;
22+
import brave.Tracing;
23+
24+
public class BraveStatementContextListener implements StatementContextListener {
25+
private final Tracing tracing;
26+
private static final Pattern serviceNamePattern = Pattern.compile("(^|&)zipkinServiceName=(?<serviceName>[^&]*)");
27+
private final Map<String, String> serviceNameCache = new ConcurrentHashMap<String, String>() { };
28+
29+
public BraveStatementContextListener(Tracing tracing) { this.tracing = tracing; }
30+
31+
@Override
32+
public void contextCreated(final StatementContext ctx) {
33+
String spanName = ctx.describeJdbiStatementType();
34+
Span span = tracing.tracer().nextSpan().name(spanName);
35+
36+
Instant start = ctx.getExecutionMoment();
37+
if (start != null) {
38+
span.start(start.toEpochMilli() * 1000);
39+
} else {
40+
span.start();
41+
}
42+
43+
ctx.setTraceId(span.context().traceIdString());
44+
45+
ctx.addCleanable(() -> {
46+
final SqlStatements stmtConfig = ctx.getConfig(SqlStatements.class);
47+
48+
String remoteServiceName = getRemoteServiceName(ctx);
49+
if (remoteServiceName != null && !remoteServiceName.isEmpty()) {
50+
span.remoteServiceName(remoteServiceName);
51+
}
52+
53+
span.kind(Span.Kind.CLIENT);
54+
55+
final String renderedSql = ctx.getRenderedSql();
56+
if (renderedSql != null) {
57+
String truncated = renderedSql.substring(
58+
0,
59+
Math.min(renderedSql.length(), stmtConfig.getJfrSqlMaxLength())
60+
);
61+
span.tag("sql.query", truncated);
62+
}
63+
64+
span.tag("sql.rows", Long.toString(ctx.getMappedRows()));
65+
66+
if (ctx.getCompletionMoment() == null) {
67+
span.tag("error", "");
68+
}
69+
70+
if (ctx.getCompletionMoment() != null) {
71+
span.finish(ctx.getCompletionMoment().toEpochMilli() * 1000);
72+
} else if (ctx.getExceptionMoment() != null) {
73+
span.finish(ctx.getExceptionMoment().toEpochMilli() * 1000);
74+
} else {
75+
span.finish();
76+
}
77+
});
78+
}
79+
80+
String getRemoteServiceName(StatementContext ctx) {
81+
Connection connection = ctx.getConnection();
82+
if (connection == null) {
83+
return null;
84+
}
85+
try {
86+
String url = connection.getMetaData().getURL();
87+
return serviceNameCache.computeIfAbsent(url, key -> {
88+
if (key.startsWith("jdbc:")) {
89+
// strip "jdbc:" prefix
90+
key = key.substring(5);
91+
}
92+
URI uri = null; // strip "jdbc:"
93+
try {
94+
uri = new URI(key);
95+
} catch (URISyntaxException ignored) {
96+
return null;
97+
}
98+
99+
String remoteServiceName = extractRemoteServiceName(uri);
100+
serviceNameCache.put(key, remoteServiceName);
101+
return remoteServiceName;
102+
});
103+
} catch (SQLException ignored) {
104+
return null;
105+
}
106+
}
107+
108+
static String extractRemoteServiceName(URI uri) {
109+
String remoteServiceNameFromQueryParameter = getRemoteServiceNameFromQueryParameter(uri);
110+
if (remoteServiceNameFromQueryParameter != null) {
111+
return remoteServiceNameFromQueryParameter;
112+
}
113+
114+
if (uri.getAuthority() != null) {
115+
return uri.getAuthority();
116+
}
117+
118+
if (uri.getScheme() != null) {
119+
return uri.getScheme();
120+
}
121+
122+
return null;
123+
}
124+
125+
static String getRemoteServiceNameFromQueryParameter(URI uri) {
126+
String query = uri.getQuery();
127+
if (query != null) {
128+
Matcher matcher = serviceNamePattern.matcher(query);
129+
if (matcher.find()) {
130+
return matcher.group("serviceName");
131+
}
132+
}
133+
134+
try {
135+
URI schemeSpecificUri = new URI(uri.getSchemeSpecificPart());
136+
if (schemeSpecificUri.toString().length() < uri.toString().length()) {
137+
return getRemoteServiceNameFromQueryParameter(schemeSpecificUri);
138+
}
139+
} catch (URISyntaxException ignored) {
140+
return null;
141+
}
142+
143+
return null;
144+
}
145+
}

‎instrumentation/jdbi3/src/main/java/brave/jdbi3/Jdbi3BravePlugin.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright The OpenZipkin Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package brave.jdbi3;
6+
7+
import java.net.URI;
8+
import java.net.URISyntaxException;
9+
import java.sql.Connection;
10+
import java.sql.SQLException;
11+
import java.util.regex.Matcher;
12+
import java.util.regex.Pattern;
13+
14+
import org.jdbi.v3.core.Jdbi;
15+
import org.jdbi.v3.core.spi.JdbiPlugin;
16+
import org.jdbi.v3.core.statement.SqlStatements;
17+
import org.jdbi.v3.core.statement.StatementContext;
18+
19+
import brave.Tracing;
20+
21+
/**
22+
* A Jdbi plugin that will report to Zipkin how long each query takes.
23+
* * <p>
24+
* Install by calling
25+
* {@code jdbi.installPlugin(Jdbi3BravePlugin.newBuilder(tracing).build());}
26+
*/
27+
public final class Jdbi3BravePlugin extends JdbiPlugin.Singleton {
28+
private final Tracing tracing;
29+
30+
private Jdbi3BravePlugin(Tracing tracing) {
31+
if (tracing == null) throw new NullPointerException("tracing == null");
32+
this.tracing = tracing;
33+
}
34+
35+
public static Builder newBuilder(Tracing tracing) {
36+
return new Builder(tracing);
37+
}
38+
39+
@Override
40+
public void customizeJdbi(Jdbi jdbi) {
41+
jdbi.getConfig(SqlStatements.class)
42+
.addContextListener(new BraveStatementContextListener(tracing));
43+
}
44+
45+
public static final class Builder {
46+
private final Tracing tracing;
47+
48+
public Builder(Tracing tracing) {
49+
this.tracing = tracing;
50+
}
51+
52+
public JdbiPlugin.Singleton build() {
53+
return new Jdbi3BravePlugin(tracing);
54+
}
55+
}
56+
}

‎instrumentation/jdbi3/src/test/java/brave/jdbi3/BraveStatementContextListenerTest.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
/*
2+
* Copyright The OpenZipkin Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package brave.jdbi3;
6+
7+
import static org.assertj.core.api.Assertions.assertThat;
8+
9+
import java.sql.Connection;
10+
import java.sql.DatabaseMetaData;
11+
import java.sql.PreparedStatement;
12+
import java.sql.SQLException;
13+
14+
import org.jdbi.v3.core.ConnectionFactory;
15+
import org.jdbi.v3.core.Jdbi;
16+
import org.jdbi.v3.core.statement.SqlStatements;
17+
import org.jdbi.v3.core.statement.StatementContext;
18+
import org.junit.jupiter.api.Test;
19+
import org.junit.jupiter.params.ParameterizedTest;
20+
import org.junit.jupiter.params.provider.CsvSource;
21+
import org.mockito.Mockito;
22+
23+
import brave.Span;
24+
import brave.Tracer;
25+
import brave.Tracing;
26+
import brave.propagation.TraceContext;
27+
28+
public class BraveStatementContextListenerTest {
29+
@Test
30+
public void shouldReportSpans() throws SQLException {
31+
Span span = Mockito.mock(Span.class);
32+
Tracing tracing = buildTracing(span);
33+
Jdbi jdbi = Jdbi.create(buildConnectionFactory());
34+
35+
jdbi.getConfig(SqlStatements.class).addContextListener(new BraveStatementContextListener(tracing));
36+
jdbi.useHandle(handle -> {
37+
handle.execute("INSERT INTO testdb.testdb (id, name) VALUES (1, 'testdb')");
38+
});
39+
40+
Mockito.verify(span).name("Update");
41+
}
42+
43+
private static ConnectionFactory buildConnectionFactory() throws SQLException {
44+
DatabaseMetaData databaseMetaData = Mockito.mock(DatabaseMetaData.class);
45+
Mockito.when(databaseMetaData.getURL()).thenReturn("jdbc:mysql://localhost:3306/testdb");
46+
Connection connection = Mockito.mock(Connection.class);
47+
Mockito.when(connection.getMetaData()).thenReturn(databaseMetaData);
48+
Mockito.when(connection.prepareStatement(Mockito.anyString(), Mockito.anyInt(), Mockito.anyInt()))
49+
.thenReturn(Mockito.mock(PreparedStatement.class));
50+
ConnectionFactory connectionFactory = () -> connection;
51+
return connectionFactory;
52+
}
53+
54+
private static Tracing buildTracing(Span span) {
55+
TraceContext traceContext = Mockito.mock(TraceContext.class);
56+
Mockito.when(span.name(Mockito.anyString())).thenReturn(span);
57+
Mockito.when(span.context()).thenReturn(traceContext);
58+
Tracer tracer = Mockito.mock(Tracer.class);
59+
Mockito.when(tracer.nextSpan()).thenReturn(span);
60+
Tracing tracing = Mockito.mock(Tracing.class);
61+
Mockito.when(tracing.tracer()).thenReturn(tracer);
62+
return tracing;
63+
}
64+
65+
@ParameterizedTest(name = "DISPLAY_NAME_PLACEHOLDER" + "{0}")
66+
@CsvSource({
67+
"jdbc:mysql://localhost:3306/testdb, localhost:3306",
68+
"jdbc:mysql://localhost:3306/testdb?zipkinServiceName=test, test",
69+
"jdbc:postgresql://localhost:5432/testdb, localhost:5432",
70+
"jdbc:postgresql://localhost:5432/testdb?zipkinServiceName=test, test",
71+
"jdbc:h2:mem:testdb, h2",
72+
"jdbc:sqlite:test.db, sqlite",
73+
"jdbc:sqlite:test.db?zipkinServiceName=test, test",
74+
"jdbc:oracle:thin:@localhost:1521:testdb, oracle",
75+
"jdbc:oracle:thin:@(DESCRIPTION=(ADDRESS=(PROTOCOL=TCP)(HOST=myhost)(PORT=1521))(CONNECT_DATA=(SERVICE_NAME=myservice))), oracle",
76+
"jdbc:oracle:oci:@//myhost.example.com:1521/MYSERVICE?zipkinServiceName=oracledb, oracledb"
77+
})
78+
public void shouldParseRemoteServiceName(String url, String expected) throws SQLException {
79+
// Mocking the StatementContext and all the others in order to be able to pass the
80+
// url as a String, and not have it parsed by the URI class etc., as that´s how
81+
// the plugin gets it from JDBC.
82+
DatabaseMetaData metadata = Mockito.mock(DatabaseMetaData.class);
83+
Mockito.when(metadata.getURL()).thenReturn(url);
84+
Connection connection = Mockito.mock(Connection.class);
85+
Mockito.when(connection.getMetaData()).thenReturn(metadata);
86+
StatementContext ctx = Mockito.mock(StatementContext.class);
87+
Mockito.when(ctx.getConnection()).thenReturn(connection);
88+
BraveStatementContextListener listener = new BraveStatementContextListener(null);
89+
String remoteServiceName = listener.getRemoteServiceName(ctx);
90+
91+
assertThat(remoteServiceName).isEqualTo(expected);
92+
}
93+
}

0 commit comments

Comments
 (0)
⚡