Skip to content

Commit 65f563f

Browse files
reftelreftel
committed
Add JDBI 3 plugin
1 parent 24e5d41 commit 65f563f

File tree

11 files changed

+529
-2
lines changed

11 files changed

+529
-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: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
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.HashMap;
13+
import java.util.Map;
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 Map<String, String> serviceNameCache = new HashMap<>();
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+
URI uri;
86+
try {
87+
String url = connection.getMetaData().getURL();
88+
if (serviceNameCache.containsKey(url)) {
89+
return serviceNameCache.get(url);
90+
}
91+
if (url.startsWith("jdbc:")) {
92+
// strip "jdbc:" prefix
93+
url = url.substring(5);
94+
}
95+
uri = new URI(url); // strip "jdbc:"
96+
97+
String remoteServiceName = extractRemoteServiceName(uri);
98+
serviceNameCache.put(url, remoteServiceName);
99+
return remoteServiceName;
100+
} catch (SQLException | URISyntaxException ignored) {
101+
return null;
102+
}
103+
}
104+
105+
static String extractRemoteServiceName(URI uri) {
106+
String remoteServiceNameFromQueryParameter = getRemoteServiceNameFromQueryParameter(uri);
107+
if (remoteServiceNameFromQueryParameter != null) {
108+
return remoteServiceNameFromQueryParameter;
109+
}
110+
111+
if (uri.getAuthority() != null) {
112+
return uri.getAuthority();
113+
}
114+
115+
if (uri.getScheme() != null) {
116+
return uri.getScheme();
117+
}
118+
119+
return null;
120+
}
121+
122+
static String getRemoteServiceNameFromQueryParameter(URI uri) {
123+
String query = uri.getQuery();
124+
if (query != null) {
125+
Matcher matcher = serviceNamePattern.matcher(query);
126+
if (matcher.find()) {
127+
return matcher.group("serviceName");
128+
}
129+
}
130+
131+
try {
132+
URI schemeSpecificUri = new URI(uri.getSchemeSpecificPart());
133+
if (schemeSpecificUri.toString().length() < uri.toString().length()) {
134+
return getRemoteServiceNameFromQueryParameter(schemeSpecificUri);
135+
}
136+
} catch (URISyntaxException ignored) {
137+
return null;
138+
}
139+
140+
return null;
141+
}
142+
}

‎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: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright The OpenZipkin Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package brave.jdbi3;
6+
7+
import java.sql.Connection;
8+
import java.sql.DatabaseMetaData;
9+
import java.sql.PreparedStatement;
10+
import java.sql.SQLException;
11+
12+
import org.jdbi.v3.core.ConnectionFactory;
13+
import org.jdbi.v3.core.Jdbi;
14+
import org.jdbi.v3.core.statement.SqlStatements;
15+
import org.junit.jupiter.api.Test;
16+
import org.mockito.Mockito;
17+
18+
import brave.Span;
19+
import brave.Tracer;
20+
import brave.Tracing;
21+
import brave.propagation.TraceContext;
22+
23+
public class BraveStatementContextListenerTest {
24+
@Test
25+
public void shouldReportSpans() throws SQLException {
26+
Span span = Mockito.mock(Span.class);
27+
Tracing tracing = buildTracing(span);
28+
Jdbi jdbi = Jdbi.create(buildConnectionFactory());
29+
30+
jdbi.getConfig(SqlStatements.class).addContextListener(new BraveStatementContextListener(tracing));
31+
jdbi.useHandle(handle -> {
32+
handle.execute("INSERT INTO testdb.testdb (id, name) VALUES (1, 'testdb')");
33+
});
34+
35+
Mockito.verify(span).name("Update");
36+
}
37+
38+
private static ConnectionFactory buildConnectionFactory() throws SQLException {
39+
DatabaseMetaData databaseMetaData = Mockito.mock(DatabaseMetaData.class);
40+
Mockito.when(databaseMetaData.getURL()).thenReturn("jdbc:mysql://localhost:3306/testdb");
41+
Connection connection = Mockito.mock(Connection.class);
42+
Mockito.when(connection.getMetaData()).thenReturn(databaseMetaData);
43+
Mockito.when(connection.prepareStatement(Mockito.anyString(), Mockito.anyInt(), Mockito.anyInt()))
44+
.thenReturn(Mockito.mock(PreparedStatement.class));
45+
ConnectionFactory connectionFactory = () -> connection;
46+
return connectionFactory;
47+
}
48+
49+
private static Tracing buildTracing(Span span) {
50+
TraceContext traceContext = Mockito.mock(TraceContext.class);
51+
Mockito.when(span.name(Mockito.anyString())).thenReturn(span);
52+
Mockito.when(span.context()).thenReturn(traceContext);
53+
Tracer tracer = Mockito.mock(Tracer.class);
54+
Mockito.when(tracer.nextSpan()).thenReturn(span);
55+
Tracing tracing = Mockito.mock(Tracing.class);
56+
Mockito.when(tracing.tracer()).thenReturn(tracer);
57+
return tracing;
58+
}
59+
}

0 commit comments

Comments
 (0)
⚡