Skip to content

Commit f756d98

Browse files
reftelreftel
committed
Add JDBI 3 plugin
1 parent bef7aba commit f756d98

File tree

5 files changed

+318
-0
lines changed

5 files changed

+318
-0
lines changed

‎instrumentation/jdbi3/pom.xml‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
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.1.1-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.0</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.h2database</groupId>
42+
<artifactId>h2</artifactId>
43+
<version>2.3.232</version>
44+
<scope>test</scope>
45+
</dependency>
46+
</dependencies>
47+
48+
<build>
49+
<plugins>
50+
<plugin>
51+
<artifactId>maven-invoker-plugin</artifactId>
52+
</plugin>
53+
<plugin>
54+
<groupId>de.qaware.maven</groupId>
55+
<artifactId>go-offline-maven-plugin</artifactId>
56+
<executions>
57+
<execution>
58+
<phase>package</phase>
59+
<goals>
60+
<goal>resolve-dependencies</goal>
61+
</goals>
62+
</execution>
63+
</executions>
64+
<configuration>
65+
<!-- Add dependencies specific to invoker tests so that they cache on go-offline -->
66+
<dynamicDependencies>
67+
<DynamicDependency>
68+
<groupId>org.jdbi</groupId>
69+
<artifactId>jdbi-3</artifactId>
70+
<version>${jdbi3.version}</version>
71+
<repositoryType>MAIN</repositoryType>
72+
<type>jar</type>
73+
</DynamicDependency>
74+
</dynamicDependencies>
75+
</configuration>
76+
</plugin>
77+
</plugins>
78+
</build>
79+
</project>

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

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
/*
2+
* Copyright The OpenZipkin Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package brave.jdbi3;
6+
7+
import java.time.Instant;
8+
9+
import org.jdbi.v3.core.Jdbi;
10+
import org.jdbi.v3.core.spi.JdbiPlugin;
11+
import org.jdbi.v3.core.statement.SqlStatements;
12+
import org.jdbi.v3.core.statement.StatementContext;
13+
import org.jdbi.v3.core.statement.StatementContextListener;
14+
15+
import brave.Span;
16+
import brave.Tracing;
17+
18+
/**
19+
* Instrumentation for JDBI 3.
20+
21+
* JDBI plugin that sends a trace span for every SQL command executed by Jdbi.
22+
* Install by calling {@code jdbi.installPlugin(new Jdbi3BravePlugin(tracing));}
23+
*/
24+
public class Jdbi3BravePlugin extends JdbiPlugin.Singleton {
25+
private final Tracing tracing;
26+
private final boolean includeBindVariables;
27+
28+
public Jdbi3BravePlugin(Tracing tracing) {
29+
this(tracing, false);
30+
}
31+
32+
public Jdbi3BravePlugin(Tracing tracing, boolean includeBindVariables) {
33+
this.tracing = tracing;
34+
this.includeBindVariables = includeBindVariables;
35+
}
36+
37+
@Override
38+
public void customizeJdbi(Jdbi jdbi) {
39+
if (tracing == null) {
40+
return;
41+
}
42+
jdbi.getConfig(SqlStatements.class)
43+
.addContextListener(new StatementContextListener() {
44+
@Override
45+
public void contextCreated(final StatementContext ctx) {
46+
String spanName = "jdbi." + ctx.describeJdbiStatementType();
47+
Span span = tracing.tracer().nextSpan().name(spanName);
48+
49+
Instant start = ctx.getExecutionMoment();
50+
if (start != null) {
51+
span.start(start.toEpochMilli() * 1000);
52+
} else {
53+
span.start();
54+
}
55+
56+
ctx.setTraceId(span.context().traceIdString());
57+
58+
ctx.addCleanable(() -> {
59+
final SqlStatements stmtConfig = ctx.getConfig(SqlStatements.class);
60+
61+
span.kind(Span.Kind.CLIENT);
62+
63+
final String renderedSql = ctx.getRenderedSql();
64+
if (renderedSql != null) {
65+
String truncated = renderedSql.substring(
66+
0,
67+
Math.min(renderedSql.length(), stmtConfig.getJfrSqlMaxLength())
68+
);
69+
span.tag("sql.query", truncated);
70+
}
71+
72+
if (includeBindVariables) {
73+
String bindVariables = ctx.getBinding()
74+
.describe(stmtConfig.getJfrParamMaxLength());
75+
span.tag("sql.binding", bindVariables);
76+
}
77+
78+
span.tag("sql.rows", Long.toString(ctx.getMappedRows()));
79+
80+
if (ctx.getCompletionMoment() != null) {
81+
span.finish(ctx.getCompletionMoment().toEpochMilli() * 1000);
82+
} else if (ctx.getExceptionMoment() != null) {
83+
span.finish(ctx.getExceptionMoment().toEpochMilli() * 1000);
84+
} else {
85+
span.finish();
86+
}
87+
});
88+
}
89+
});
90+
}
91+
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
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.SQLException;
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
import java.util.Map;
13+
14+
import org.assertj.core.api.Assertions;
15+
import org.jdbi.v3.core.Jdbi;
16+
import org.jdbi.v3.core.statement.Cleanable;
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+
import org.junit.jupiter.api.Test;
21+
import org.mockito.ArgumentCaptor;
22+
import org.mockito.Mockito;
23+
24+
import brave.Span;
25+
import brave.Tracer;
26+
import brave.Tracing;
27+
import brave.handler.MutableSpan;
28+
import brave.handler.SpanHandler;
29+
import brave.propagation.TraceContext;
30+
31+
class Jdbi3BravePluginIT {
32+
@Test
33+
public void shouldSendSpans() {
34+
ListSpanHandler spanHandler = new ListSpanHandler();
35+
try (
36+
Tracing tracing = Tracing.newBuilder()
37+
.addSpanHandler(spanHandler)
38+
.build()
39+
) {
40+
Jdbi jdbi = Jdbi.create("jdbc:h2:mem:test");
41+
jdbi.installPlugin(new Jdbi3BravePlugin(tracing, true));
42+
jdbi.withHandle(h -> {
43+
h.execute("CREATE TABLE test (id INT PRIMARY KEY, name VARCHAR(255))");
44+
h.execute("INSERT INTO test (id, name) VALUES (1, 'test')");
45+
return h.createQuery("SELECT * FROM test where name = :name")
46+
.bind("name", "test")
47+
.mapToMap()
48+
.list();
49+
});
50+
}
51+
52+
assertThat(spanHandler.getSpans())
53+
.isNotEmpty()
54+
.anySatisfy(span -> {
55+
assertThat(span.kind()).isEqualTo(Span.Kind.CLIENT);
56+
assertThat(span.name()).isEqualTo("jdbi.Query");
57+
assertThat(span.tags()).containsEntry("sql.query", "SELECT * FROM test where name = :name")
58+
.containsEntry("sql.rows", "1")
59+
.containsEntry("sql.binding", "{named:{name:test}}");
60+
});
61+
}
62+
63+
private static class ListSpanHandler extends SpanHandler {
64+
private final List<MutableSpan> spans = new ArrayList<>();
65+
66+
@Override
67+
public boolean end(TraceContext context, MutableSpan span, Cause cause) {
68+
spans.add(span);
69+
return super.end(context, span, cause);
70+
}
71+
72+
public List<MutableSpan> getSpans() {
73+
return spans;
74+
}
75+
}
76+
}

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*
2+
* Copyright The OpenZipkin Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package brave.jdbi3;
6+
7+
import static org.junit.jupiter.api.Assertions.*;
8+
9+
import java.sql.SQLException;
10+
11+
import org.jdbi.v3.core.Jdbi;
12+
import org.jdbi.v3.core.statement.Cleanable;
13+
import org.jdbi.v3.core.statement.SqlStatements;
14+
import org.jdbi.v3.core.statement.StatementContext;
15+
import org.jdbi.v3.core.statement.StatementContextListener;
16+
import org.junit.jupiter.api.Test;
17+
import org.mockito.ArgumentCaptor;
18+
import org.mockito.Mockito;
19+
20+
import brave.Span;
21+
import brave.Tracer;
22+
import brave.Tracing;
23+
import brave.propagation.TraceContext;
24+
25+
class Jdbi3BravePluginTest {
26+
@Test
27+
public void shouldSendSpans() throws SQLException {
28+
Span span = Mockito.mock(Span.class);
29+
Tracer tracer = Mockito.mock(Tracer.class);
30+
Mockito.when(tracer.nextSpan()).thenReturn(span);
31+
Mockito.when(span.name(Mockito.anyString())).thenReturn(span);
32+
TraceContext traceContext = Mockito.mock(TraceContext.class);
33+
Mockito.when(span.context()).thenReturn(traceContext);
34+
Tracing tracing = Mockito.mock(Tracing.class);
35+
Mockito.when(tracing.tracer()).thenReturn(tracer);
36+
Jdbi jdbi = Mockito.mock(Jdbi.class);
37+
SqlStatements sqlStatements = Mockito.mock(SqlStatements.class);
38+
Mockito.when(sqlStatements.addContextListener(Mockito.any())).thenReturn(sqlStatements);
39+
Mockito.when(sqlStatements.getJfrSqlMaxLength()).thenReturn(100);
40+
Mockito.when(jdbi.getConfig(SqlStatements.class)).thenReturn(sqlStatements);
41+
42+
Jdbi3BravePlugin plugin = new Jdbi3BravePlugin(tracing);
43+
44+
plugin.customizeJdbi(jdbi);
45+
46+
ArgumentCaptor<StatementContextListener> statementContextListenerArgumentCaptor =
47+
ArgumentCaptor.forClass(StatementContextListener.class);
48+
Mockito.verify(sqlStatements, Mockito.times(1))
49+
.addContextListener(statementContextListenerArgumentCaptor.capture());
50+
51+
StatementContext statementContext = Mockito.mock(StatementContext.class);
52+
Mockito.when(statementContext.describeJdbiStatementType()).thenReturn("statement");
53+
Mockito.when(statementContext.getRenderedSql()).thenReturn("SELECT * FROM table");
54+
Mockito.when(statementContext.getConfig(SqlStatements.class)).thenReturn(sqlStatements);
55+
Mockito.when(statementContext.getMappedRows()).thenReturn(42L);
56+
57+
statementContextListenerArgumentCaptor.getValue().contextCreated(statementContext);
58+
59+
ArgumentCaptor<Cleanable> cleanableArgumentCaptor = ArgumentCaptor.forClass(Cleanable.class);
60+
Mockito.verify(statementContext, Mockito.times(1)).addCleanable(cleanableArgumentCaptor.capture());
61+
62+
cleanableArgumentCaptor.getValue().close();
63+
64+
Mockito.verify(span, Mockito.times(1)).start();
65+
Mockito.verify(span, Mockito.times(1)).name("jdbi.statement");
66+
Mockito.verify(span, Mockito.times(1)).kind(Span.Kind.CLIENT);
67+
Mockito.verify(span, Mockito.times(1)).tag("sql.query", "SELECT * FROM table");
68+
Mockito.verify(span, Mockito.times(1)).tag("sql.rows", "42");
69+
Mockito.verify(span, Mockito.times(1)).finish();
70+
}
71+
}

‎instrumentation/pom.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
<module>httpclient5</module>
3434
<module>jaxrs2</module>
3535
<module>jersey-server</module>
36+
<module>jdbi3</module>
3637
<module>jms</module>
3738
<module>jms-jakarta</module>
3839
<module>jakarta-jms</module>

0 commit comments

Comments
 (0)
⚡