Skip to content

Commit 113167e

Browse files
eyalkoreneyalkoren
authored
Add support for JDBC executeBatch (#747)
1 parent b796ce2 commit 113167e

File tree

8 files changed

+182
-177
lines changed

8 files changed

+182
-177
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
## Bug Fixes
99
* Some JMS Consumers and Producers are filtered due to class name filtering in instrumentation matching
1010
* Jetty: When no display name is set and context path is "/" transaction service names will now correctly fall back to configured values
11+
* JDBC's `executeBatch` is not traced
1112
* Drops non-String labels when connected to APM Server < 6.7 to avoid validation errors (#687)
1213

1314
# 1.7.0

‎apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/ConnectionInstrumentation.java‎

Lines changed: 2 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,13 @@
2626

2727
import co.elastic.apm.agent.bci.ElasticApmInstrumentation;
2828
import co.elastic.apm.agent.bci.VisibleForAdvice;
29-
import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
29+
import co.elastic.apm.agent.jdbc.helper.JdbcHelper;
3030
import net.bytebuddy.asm.Advice;
3131
import net.bytebuddy.description.NamedElement;
3232
import net.bytebuddy.description.method.MethodDescription;
3333
import net.bytebuddy.description.type.TypeDescription;
3434
import net.bytebuddy.matcher.ElementMatcher;
3535

36-
import javax.annotation.Nullable;
3736
import java.sql.Connection;
3837
import java.sql.PreparedStatement;
3938
import java.util.Collection;
@@ -55,28 +54,12 @@
5554
*/
5655
public class ConnectionInstrumentation extends ElasticApmInstrumentation {
5756

58-
@VisibleForAdvice
59-
public static final WeakConcurrentMap<Object, String> statementSqlMap = new WeakConcurrentMap<Object, String>(true);
6057
static final String JDBC_INSTRUMENTATION_GROUP = "jdbc";
6158

6259
@VisibleForAdvice
6360
@Advice.OnMethodExit(suppress = Throwable.class)
6461
public static void storeSql(@Advice.Return final PreparedStatement statement, @Advice.Argument(0) String sql) {
65-
statementSqlMap.putIfAbsent(statement, sql);
66-
}
67-
68-
/**
69-
* Returns the SQL statement belonging to provided {@link PreparedStatement}.
70-
* <p>
71-
* Might return {@code null} when the provided {@link PreparedStatement} is a wrapper of the actual statement.
72-
* </p>
73-
*
74-
* @return the SQL statement belonging to provided {@link PreparedStatement}, or {@code null}
75-
*/
76-
@Nullable
77-
@VisibleForAdvice
78-
public static String getSqlForStatement(Object statement) {
79-
return statementSqlMap.get(statement);
62+
JdbcHelper.mapStatementToSql(statement, sql);
8063
}
8164

8265
@Override

‎apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/PreparedStatementInstrumentation.java‎

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

‎apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/StatementInstrumentation.java‎

Lines changed: 97 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -51,45 +51,28 @@
5151
import static net.bytebuddy.matcher.ElementMatchers.named;
5252
import static net.bytebuddy.matcher.ElementMatchers.not;
5353
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
54+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
5455

5556
/**
5657
* Creates spans for JDBC calls
5758
*/
58-
public class StatementInstrumentation extends ElasticApmInstrumentation {
59+
public abstract class StatementInstrumentation extends ElasticApmInstrumentation {
5960

61+
@SuppressWarnings("WeakerAccess")
6062
@Nullable
6163
@VisibleForAdvice
6264
public static HelperClassManager<JdbcHelper> jdbcHelperManager;
6365

64-
public StatementInstrumentation(ElasticApmTracer tracer) {
65-
jdbcHelperManager = HelperClassManager.ForSingleClassLoader.of(tracer, "co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl",
66+
private final ElementMatcher<? super MethodDescription> methodMatcher;
67+
68+
StatementInstrumentation(ElasticApmTracer tracer, ElementMatcher<? super MethodDescription> methodMatcher) {
69+
this.methodMatcher = methodMatcher;
70+
jdbcHelperManager = HelperClassManager.ForSingleClassLoader.of(tracer,
71+
"co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl",
6672
"co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl$1",
6773
"co.elastic.apm.agent.jdbc.helper.JdbcHelperImpl$ConnectionMetaData");
6874
}
6975

70-
@Nullable
71-
@VisibleForAdvice
72-
@Advice.OnMethodEnter(suppress = Throwable.class)
73-
public static Span onBeforeExecute(@Advice.This Statement statement, @Advice.Argument(0) String sql) throws SQLException {
74-
if (tracer != null && jdbcHelperManager != null) {
75-
JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
76-
if (helperImpl != null) {
77-
return helperImpl.createJdbcSpan(sql, statement.getConnection(), tracer.getActive(), false);
78-
}
79-
}
80-
return null;
81-
}
82-
83-
@VisibleForAdvice
84-
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
85-
public static void onAfterExecute(@Advice.Enter @Nullable Span span, @Advice.Thrown Throwable t) {
86-
if (span != null) {
87-
span.captureException(t)
88-
.deactivate()
89-
.end();
90-
}
91-
}
92-
9376
@Override
9477
public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
9578
return nameContains("Statement");
@@ -103,13 +86,98 @@ public ElementMatcher<? super TypeDescription> getTypeMatcher() {
10386

10487
@Override
10588
public ElementMatcher<? super MethodDescription> getMethodMatcher() {
106-
return nameStartsWith("execute")
107-
.and(takesArgument(0, String.class))
108-
.and(isPublic());
89+
return methodMatcher;
10990
}
11091

11192
@Override
11293
public Collection<String> getInstrumentationGroupNames() {
11394
return Collections.singleton(JDBC_INSTRUMENTATION_GROUP);
11495
}
96+
97+
public static class ExecuteWithQueryInstrumentation extends StatementInstrumentation {
98+
99+
public ExecuteWithQueryInstrumentation(ElasticApmTracer tracer) {
100+
super(tracer,
101+
nameStartsWith("execute")
102+
.and(takesArgument(0, String.class))
103+
.and(isPublic())
104+
);
105+
}
106+
107+
@Nullable
108+
@VisibleForAdvice
109+
@Advice.OnMethodEnter(suppress = Throwable.class)
110+
public static Span onBeforeExecute(@Advice.This Statement statement, @Advice.Argument(0) String sql) throws SQLException {
111+
if (tracer != null && jdbcHelperManager != null) {
112+
JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
113+
if (helperImpl != null) {
114+
return helperImpl.createJdbcSpan(sql, statement.getConnection(), tracer.getActive(), false);
115+
}
116+
}
117+
return null;
118+
}
119+
120+
@VisibleForAdvice
121+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
122+
public static void onAfterExecute(@Advice.Enter @Nullable Span span, @Advice.Thrown Throwable t) {
123+
if (span != null) {
124+
span.captureException(t)
125+
.deactivate()
126+
.end();
127+
}
128+
}
129+
}
130+
131+
public static class AddBatchInstrumentation extends StatementInstrumentation {
132+
133+
public AddBatchInstrumentation(ElasticApmTracer tracer) {
134+
super(tracer,
135+
nameStartsWith("addBatch")
136+
.and(takesArgument(0, String.class))
137+
.and(isPublic())
138+
);
139+
}
140+
141+
@Advice.OnMethodEnter(suppress = Throwable.class)
142+
public static void storeSql(@Advice.This Statement statement, @Advice.Argument(0) String sql) {
143+
if (jdbcHelperManager != null) {
144+
JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
145+
if (helperImpl != null) {
146+
helperImpl.mapStatementToSql(statement, sql);
147+
}
148+
}
149+
}
150+
}
151+
152+
public static class ExecuteWithoutQueryInstrumentation extends StatementInstrumentation {
153+
public ExecuteWithoutQueryInstrumentation(ElasticApmTracer tracer) {
154+
super(tracer,
155+
nameStartsWith("execute")
156+
.and(takesArguments(0))
157+
.and(isPublic())
158+
);
159+
}
160+
161+
@Nullable
162+
@Advice.OnMethodEnter(suppress = Throwable.class)
163+
public static Span onBeforeExecute(@Advice.This Statement statement) throws SQLException {
164+
if (tracer != null && jdbcHelperManager != null) {
165+
JdbcHelper helperImpl = jdbcHelperManager.getForClassLoaderOfClass(Statement.class);
166+
if (helperImpl != null) {
167+
final @Nullable String sql = helperImpl.retrieveSqlForStatement(statement);
168+
return helperImpl.createJdbcSpan(sql, statement.getConnection(), tracer.getActive(), true);
169+
}
170+
}
171+
return null;
172+
}
173+
174+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
175+
public static void onAfterExecute(@Advice.Enter @Nullable Span span, @Advice.Thrown Throwable t) {
176+
if (span != null) {
177+
span.captureException(t)
178+
.deactivate()
179+
.end();
180+
}
181+
}
182+
}
115183
}

‎apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/helper/JdbcHelper.java‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,43 @@
2424
*/
2525
package co.elastic.apm.agent.jdbc.helper;
2626

27+
import co.elastic.apm.agent.bci.VisibleForAdvice;
2728
import co.elastic.apm.agent.impl.transaction.Span;
2829
import co.elastic.apm.agent.impl.transaction.TraceContextHolder;
30+
import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
2931

3032
import javax.annotation.Nullable;
3133
import java.sql.Connection;
3234

33-
public interface JdbcHelper {
35+
public abstract class JdbcHelper {
36+
@SuppressWarnings("WeakerAccess")
37+
@VisibleForAdvice
38+
public static final WeakConcurrentMap<Object, String> statementSqlMap = new WeakConcurrentMap<>(true);
39+
40+
/**
41+
* Maps the provided sql to the provided Statement object
42+
*
43+
* @param statement javax.sql.Statement object
44+
* @param sql query string
45+
*/
46+
public static void mapStatementToSql(Object statement, String sql) {
47+
statementSqlMap.putIfAbsent(statement, sql);
48+
}
49+
50+
/**
51+
* Returns the SQL statement belonging to provided Statement.
52+
* <p>
53+
* Might return {@code null} when the provided Statement is a wrapper of the actual statement.
54+
* </p>
55+
*
56+
* @return the SQL statement belonging to provided Statement, or {@code null}
57+
*/
58+
@Nullable
59+
public static String retrieveSqlForStatement(Object statement) {
60+
return statementSqlMap.get(statement);
61+
}
62+
63+
3464
@Nullable
35-
Span createJdbcSpan(@Nullable String sql, Connection connection, @Nullable TraceContextHolder<?> parent, boolean preparedStatement);
65+
public abstract Span createJdbcSpan(@Nullable String sql, Connection connection, @Nullable TraceContextHolder<?> parent, boolean preparedStatement);
3666
}

‎apm-agent-plugins/apm-jdbc-plugin/src/main/java/co/elastic/apm/agent/jdbc/helper/JdbcHelperImpl.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
import java.sql.DatabaseMetaData;
3838
import java.sql.SQLException;
3939

40-
public class JdbcHelperImpl implements JdbcHelper {
40+
public class JdbcHelperImpl extends JdbcHelper {
4141
public static final String DB_SPAN_TYPE = "db";
4242
public static final String DB_SPAN_ACTION = "query";
4343

‎apm-agent-plugins/apm-jdbc-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.bci.ElasticApmInstrumentation‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
co.elastic.apm.agent.jdbc.ConnectionInstrumentation
2-
co.elastic.apm.agent.jdbc.StatementInstrumentation
3-
co.elastic.apm.agent.jdbc.PreparedStatementInstrumentation
2+
co.elastic.apm.agent.jdbc.StatementInstrumentation$ExecuteWithQueryInstrumentation
3+
co.elastic.apm.agent.jdbc.StatementInstrumentation$AddBatchInstrumentation
4+
co.elastic.apm.agent.jdbc.StatementInstrumentation$ExecuteWithoutQueryInstrumentation

0 commit comments

Comments
 (0)