execute = this.databaseClient.execute("INSERT INTO test VALUES (:value)")
96 | .bind("value", 100)
97 | .fetch().rowsUpdated();
98 | return this.operator.execute(status -> {
99 | status.setRollbackOnly();
100 | return execute;
101 | }).ofType(Integer.class);
102 | }
103 |
104 | // Batch
105 |
106 | // Error
107 | @GetMapping("/error")
108 | Flux error() {
109 | return this.databaseClient.execute("SELECT SOMETHING_WRONG();").map(row -> row.get(0)).all();
110 | }
111 |
112 | // Error with recovery
113 | @GetMapping("/error-recovery")
114 | Flux errorRecovery() {
115 | return this.databaseClient.execute("SELECT SOMETHING_WRONG();")
116 | .map(row -> row.get(0))
117 | .all()
118 | .onErrorReturn("recovered");
119 | }
120 |
121 | // Error with transaction
122 | @GetMapping("/error-tx")
123 | Flux errorTx() {
124 | return this.databaseClient.execute("SELECT SOMETHING_WRONG();")
125 | .map(row -> row.get(0))
126 | .all()
127 | .as(this.operator::transactional);
128 | }
129 |
130 | }
131 |
--------------------------------------------------------------------------------
/listener-example/src/main/java/io/r2dbc/examples/MetricsExecutionListener.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import io.micrometer.core.instrument.Counter;
4 | import io.micrometer.core.instrument.MeterRegistry;
5 | import io.micrometer.core.instrument.Timer;
6 | import io.r2dbc.proxy.core.MethodExecutionInfo;
7 | import io.r2dbc.proxy.core.QueryExecutionInfo;
8 | import io.r2dbc.proxy.listener.ProxyMethodExecutionListener;
9 | import io.r2dbc.proxy.support.QueryExecutionInfoFormatter;
10 | import org.slf4j.Logger;
11 | import org.slf4j.LoggerFactory;
12 |
13 | import java.time.Duration;
14 |
15 | /**
16 | * Listener to populate micrometer metrics and logs slow query.
17 | *
18 | * @author Tadaya Tsuyukubo
19 | */
20 | public class MetricsExecutionListener implements ProxyMethodExecutionListener {
21 | private static final Logger logger = LoggerFactory.getLogger(MetricsExecutionListener.class);
22 |
23 | private MeterRegistry registry;
24 |
25 | private String metricNamePrefix = "r2dbc.";
26 |
27 | private Duration slowQueryThreshold = Duration.ofSeconds(-1); // negative won't match any query
28 |
29 | private QueryExecutionInfoFormatter queryFormatter = new QueryExecutionInfoFormatter()
30 | .showTime()
31 | .showConnection()
32 | .showQuery();
33 |
34 | public MetricsExecutionListener(MeterRegistry registry) {
35 | this.registry = registry;
36 | }
37 |
38 | public MetricsExecutionListener(MeterRegistry registry, Duration slowQueryThreshold) {
39 | this.registry = registry;
40 | this.slowQueryThreshold = slowQueryThreshold;
41 | }
42 |
43 | @Override
44 | public void beforeCreateOnConnectionFactory(MethodExecutionInfo methodExecutionInfo) {
45 | Timer.Sample sample = Timer.start(this.registry);
46 | methodExecutionInfo.getValueStore().put("connectionCreate", sample);
47 | }
48 |
49 | @Override
50 | public void afterCreateOnConnectionFactory(MethodExecutionInfo methodExecutionInfo) {
51 | Timer.Sample sample = methodExecutionInfo.getValueStore().get("connectionCreate", Timer.Sample.class);
52 |
53 | Timer timer = Timer
54 | .builder(this.metricNamePrefix + "connection")
55 | .description("Time to create(acquire) a connection")
56 | .tags("event", "create")
57 | .register(this.registry);
58 |
59 | sample.stop(timer);
60 | }
61 |
62 | @Override
63 | public void afterCommitTransactionOnConnection(MethodExecutionInfo methodExecutionInfo) {
64 | Counter counter = Counter
65 | .builder(this.metricNamePrefix + "transaction")
66 | .description("Num of transactions")
67 | .tags("event", "commit")
68 | .register(registry);
69 | counter.increment();
70 | }
71 |
72 | @Override
73 | public void afterRollbackTransactionOnConnection(MethodExecutionInfo methodExecutionInfo) {
74 | incrementRollbackCounter();
75 | }
76 |
77 | @Override
78 | public void afterRollbackTransactionToSavepointOnConnection(MethodExecutionInfo methodExecutionInfo) {
79 | incrementRollbackCounter();
80 | }
81 |
82 | private void incrementRollbackCounter() {
83 | Counter counter = Counter
84 | .builder(this.metricNamePrefix + "transaction")
85 | .description("Num of transactions")
86 | .tags("event", "rollback")
87 | .register(registry);
88 | counter.increment();
89 | }
90 |
91 |
92 | @Override
93 | public void afterExecuteOnBatch(QueryExecutionInfo queryExecutionInfo) {
94 | afterExecuteQuery(queryExecutionInfo);
95 | }
96 |
97 | @Override
98 | public void afterExecuteOnStatement(QueryExecutionInfo queryExecutionInfo) {
99 | afterExecuteQuery(queryExecutionInfo);
100 | }
101 |
102 | private void afterExecuteQuery(QueryExecutionInfo queryExecutionInfo) {
103 | Counter success = Counter
104 | .builder(this.metricNamePrefix + "query")
105 | .description("Num of executed queries")
106 | .register(this.registry);
107 | success.increment();
108 |
109 |
110 | // when negative value is specified, do not log slow query
111 | if (this.slowQueryThreshold.isNegative()) {
112 | return;
113 | }
114 |
115 | if (this.slowQueryThreshold.minus(queryExecutionInfo.getExecuteDuration()).isNegative()) {
116 | Counter slowQueryCounter = Counter
117 | .builder(this.metricNamePrefix + "query.slow")
118 | .description("Slow query count that took more than threshold")
119 | .register(registry);
120 | slowQueryCounter.increment();
121 |
122 |
123 | StringBuilder sb = new StringBuilder();
124 | sb.append("SlowQuery: ");
125 | sb.append(this.queryFormatter.format(queryExecutionInfo));
126 | logger.info(sb.toString());
127 | }
128 | }
129 |
130 | public void setRegistry(MeterRegistry registry) {
131 | this.registry = registry;
132 | }
133 |
134 | public void setMetricNamePrefix(String metricNamePrefix) {
135 | this.metricNamePrefix = metricNamePrefix;
136 | }
137 |
138 | public void setSlowQueryThreshold(Duration slowQueryThreshold) {
139 | this.slowQueryThreshold = slowQueryThreshold;
140 | }
141 | }
142 |
--------------------------------------------------------------------------------
/listener-example/src/main/java/io/r2dbc/examples/QueryTimeMetricsExecutionListener.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import io.micrometer.core.instrument.MeterRegistry;
4 | import io.micrometer.core.instrument.Timer;
5 | import io.r2dbc.proxy.core.QueryExecutionInfo;
6 | import io.r2dbc.proxy.core.QueryInfo;
7 | import io.r2dbc.proxy.listener.ProxyExecutionListener;
8 |
9 | import static java.lang.String.format;
10 |
11 | /**
12 | * Listener to populate micrometer metrics for query execution.
13 | *
14 | * Create time metrics for query execution by type(read/write).
15 | * https://github.com/micrometer-metrics/micrometer/issues/635
16 | *
17 | * @author Tadaya Tsuyukubo
18 | */
19 | public class QueryTimeMetricsExecutionListener implements ProxyExecutionListener {
20 |
21 | private MeterRegistry registry;
22 |
23 | private String metricNamePrefix = "r2dbc.";
24 |
25 | private QueryTypeDetector queryTypeDetector = new DefaultQueryTypeDetector();
26 |
27 | public QueryTimeMetricsExecutionListener(MeterRegistry registry) {
28 | this.registry = registry;
29 | }
30 |
31 | @Override
32 | public void afterQuery(QueryExecutionInfo queryExecutionInfo) {
33 | for (QueryInfo queryInfo : queryExecutionInfo.getQueries()) {
34 | String queryType = this.queryTypeDetector.detect(queryInfo.getQuery()).name().toLowerCase();
35 | String metricsName = this.metricNamePrefix + "query." + queryType;
36 | String description = format("Time to execute %s queries", queryType);
37 |
38 | Timer timer = Timer
39 | .builder(metricsName)
40 | .description(description)
41 | .tags("event", "query")
42 | .register(this.registry);
43 | timer.record(queryExecutionInfo.getExecuteDuration());
44 | }
45 | }
46 |
47 |
48 | public void setRegistry(MeterRegistry registry) {
49 | this.registry = registry;
50 | }
51 |
52 | public void setMetricNamePrefix(String metricNamePrefix) {
53 | this.metricNamePrefix = metricNamePrefix;
54 | }
55 |
56 | public void setQueryTypeDetector(QueryTypeDetector queryTypeDetector) {
57 | this.queryTypeDetector = queryTypeDetector;
58 | }
59 |
60 | public enum QueryType {
61 | SELECT, INSERT, UPDATE, DELETE, OTHER
62 | }
63 |
64 | public interface QueryTypeDetector {
65 | QueryType detect(String query);
66 | }
67 |
68 | public static class DefaultQueryTypeDetector implements QueryTypeDetector {
69 | @Override
70 | public QueryType detect(String query) {
71 | final String trimmedQuery = removeCommentAndWhiteSpace(query);
72 | if (trimmedQuery == null || trimmedQuery.length() < 6) {
73 | return QueryType.OTHER;
74 | }
75 |
76 | String prefix = trimmedQuery.substring(0, 6).toUpperCase();
77 | final QueryType type;
78 | switch (prefix) {
79 | case "SELECT":
80 | type = QueryType.SELECT;
81 | break;
82 | case "INSERT":
83 | type = QueryType.INSERT;
84 | break;
85 | case "UPDATE":
86 | type = QueryType.UPDATE;
87 | break;
88 | case "DELETE":
89 | type = QueryType.DELETE;
90 | break;
91 | default:
92 | type = QueryType.OTHER;
93 | }
94 | return type;
95 | }
96 |
97 | private String removeCommentAndWhiteSpace(String query) {
98 | if (query == null) {
99 | return null;
100 | }
101 | return query.replaceAll("--.*\n", "").replaceAll("\n", "").replaceAll("/\\*.*\\*/", "").trim();
102 | }
103 |
104 | }
105 |
106 |
107 | }
108 |
--------------------------------------------------------------------------------
/listener-example/src/main/java/io/r2dbc/examples/R2dbcSpiController.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import io.r2dbc.spi.Connection;
4 | import io.r2dbc.spi.ConnectionFactory;
5 | import io.r2dbc.spi.Result;
6 | import org.springframework.web.bind.annotation.GetMapping;
7 | import org.springframework.web.bind.annotation.RequestMapping;
8 | import org.springframework.web.bind.annotation.RestController;
9 | import reactor.core.publisher.Flux;
10 | import reactor.core.publisher.Mono;
11 |
12 | import java.util.function.Function;
13 | import java.util.stream.Collectors;
14 |
15 | /**
16 | * Provide endpoints using R2DBC SPIs.
17 | *
18 | * @author Tadaya Tsuyukubo
19 | */
20 | @RestController
21 | @RequestMapping(path = "/spi")
22 | public class R2dbcSpiController {
23 |
24 | private final ConnectionFactory connectionFactory;
25 |
26 | public R2dbcSpiController(ConnectionFactory connectionFactory) {
27 | this.connectionFactory = connectionFactory;
28 | }
29 |
30 | @GetMapping
31 | String hi() {
32 | return "Hi from " + getClass().getSimpleName();
33 | }
34 |
35 | // Perform a query without transaction
36 | @GetMapping("/simple")
37 | Flux simple() {
38 | return Flux.usingWhen(this.connectionFactory.create(), connection -> {
39 | String query = "SELECT value FROM test";
40 | Flux execute = Flux.from(connection.createStatement(query).execute());
41 | Function> mapper = (result) -> Flux.from(result.map((row, rowMetadata) -> row.get("value", Integer.class)));
42 | return execute.flatMap(mapper);
43 | }, Connection::close, (c, err) -> c.close(), Connection::close);
44 | }
45 |
46 | @GetMapping("/multi-queries")
47 | Flux multiQueries() {
48 | return Flux.usingWhen(this.connectionFactory.create(), connection -> {
49 | String queries = "SELECT value FROM test; SELECT value FROM test";
50 | Flux execute = Flux.from(connection.createStatement(queries).execute());
51 | Function> mapper = (result) -> Flux.from(result.map((row, rowMetadata) -> row.get("value", Integer.class)));
52 | return execute.flatMap(mapper);
53 | }, Connection::close, (c, err) -> c.close(), Connection::close);
54 | }
55 |
56 | // Perform an update query with transaction
57 | @GetMapping("/tx")
58 | Mono tx() {
59 | return Mono.usingWhen(this.connectionFactory.create(), connection -> {
60 | String query = "INSERT INTO test VALUES ($1)";
61 | Flux execute = Flux.from(connection.createStatement(query).bind("$1", "100").execute());
62 | Function> mapper = (result) -> Mono.from(result.getRowsUpdated());
63 | Flux action = execute.flatMap(mapper);
64 |
65 | return Flux.concat(
66 | Mono.from(connection.beginTransaction()).then(Mono.empty()),
67 | action,
68 | Mono.from(connection.commitTransaction()).then(Mono.empty())
69 | ).collect(Collectors.summingInt(i -> i));
70 | }, Connection::close, (c, err) -> c.close(), Connection::close);
71 | }
72 |
73 | @GetMapping("/tx-with-queries")
74 | Mono txWithQueries() {
75 | return Mono.usingWhen(this.connectionFactory.create(), connection -> {
76 | String query = "INSERT INTO test VALUES ($1)";
77 | Function> mapper = (result) -> Mono.from(result.getRowsUpdated());
78 | Flux execute1 = Flux.from(connection.createStatement(query).bind("$1", "100").execute()).flatMap(mapper);
79 | Flux execute2 = Flux.from(connection.createStatement(query).bind("$1", "200").execute()).flatMap(mapper);
80 |
81 | // add up num of updated rows
82 | Mono action = Flux.concat(execute1, execute2).collect(Collectors.summingInt(value -> value));
83 |
84 | return Flux.concat(
85 | Mono.from(connection.beginTransaction()).then(Mono.empty()),
86 | action,
87 | Mono.from(connection.commitTransaction()).then(Mono.empty())
88 | ).collect(Collectors.summingInt(i -> i));
89 | }, Connection::close, (c, err) -> c.close(), Connection::close);
90 | }
91 |
92 |
93 | // Multiple Tx on single connection
94 | @GetMapping("/multi-tx")
95 | Flux multipleTx() {
96 | return Flux.usingWhen(this.connectionFactory.create(), connection -> {
97 | String query = "INSERT INTO test VALUES ($1)";
98 | Function> mapper = (result) -> Mono.from(result.getRowsUpdated());
99 | Flux execute1 = Flux.from(connection.createStatement(query).bind("$1", "100").execute()).flatMap(mapper);
100 | Flux execute2 = Flux.from(connection.createStatement(query).bind("$1", "200").execute()).flatMap(mapper);
101 |
102 | return Flux.concat(
103 | Mono.from(connection.beginTransaction()).then(Mono.empty()),
104 | execute1,
105 | Mono.from(connection.commitTransaction()).then(Mono.empty()),
106 | Mono.from(connection.beginTransaction()).then(Mono.empty()),
107 | execute2,
108 | Mono.from(connection.commitTransaction()).then(Mono.empty())
109 | );
110 | }, Connection::close, (c, err) -> c.close(), Connection::close);
111 | }
112 |
113 | // Explicit rollback
114 | @GetMapping("/rollback")
115 | Mono rollback() {
116 | return Mono.usingWhen(this.connectionFactory.create(), connection -> {
117 | String query = "INSERT INTO test VALUES ($1)";
118 | Function> mapper = (result) -> Mono.from(result.getRowsUpdated());
119 | Flux execute = Flux.from(connection.createStatement(query).bind("$1", "100").execute()).flatMap(mapper);
120 |
121 | return Flux.concat(
122 | Mono.from(connection.beginTransaction()).then(Mono.empty()),
123 | execute,
124 | Mono.from(connection.rollbackTransaction()).then(Mono.empty())
125 | ).collect(Collectors.summingInt(i -> i));
126 | }, Connection::close, (c, err) -> c.close(), Connection::close);
127 | }
128 |
129 | // Batch
130 | @GetMapping("/batch")
131 | Flux batch() {
132 | return Flux.usingWhen(this.connectionFactory.create(), connection -> {
133 | String query1 = "INSERT INTO test VALUES (50)";
134 | String query2 = "INSERT INTO test VALUES (70)";
135 | Function> mapper = (result) -> Mono.from(result.getRowsUpdated());
136 | return Flux.from(connection.createBatch().add(query1).add(query2).execute()).flatMap(mapper);
137 | }, Connection::close, (c, err) -> c.close(), Connection::close);
138 | }
139 |
140 | // Error
141 | @GetMapping("/error")
142 | Flux error() {
143 | return Flux.usingWhen(this.connectionFactory.create(), connection -> {
144 | String query = "SELECT SOMETHING_WRONG();";
145 | Flux execute = Flux.from(connection.createStatement(query).execute());
146 | Function> mapper = (result) -> Flux.from(result.map((row, rowMetadata) -> row.get(0, Integer.class)));
147 | return execute.flatMap(mapper);
148 | }, Connection::close, (c, err) -> c.close(), Connection::close);
149 | }
150 |
151 | // Error with recovery
152 | @GetMapping("/error-recovery")
153 | Flux errorRecovery() {
154 | return Flux.usingWhen(this.connectionFactory.create(), connection -> {
155 | String query = "SELECT SOMETHING_WRONG();";
156 | Flux execute = Flux.from(connection.createStatement(query).execute());
157 | Function> mapper = (result) -> Flux.from(result.map((row, rowMetadata) -> row.get(0, Integer.class)));
158 | return execute.flatMap(mapper).onErrorReturn(-1);
159 | }, Connection::close, (c, err) -> c.close(), Connection::close);
160 | }
161 |
162 | // Error with Tx
163 | @GetMapping("/error-tx")
164 | Mono errorTx() {
165 | return Mono.usingWhen(this.connectionFactory.create(), connection -> {
166 | String query = "INSERT INTO test VALUES ($1)";
167 | Flux execute = Flux.from(connection.createStatement(query).bind("$1", "ABC").execute());
168 | Function> mapper = (result) -> Mono.from(result.getRowsUpdated());
169 | Flux action = execute.flatMap(mapper);
170 |
171 | return Flux.concat(
172 | Mono.from(connection.beginTransaction()).then(),
173 | action,
174 | Mono.from(connection.commitTransaction()).then()
175 | ).onErrorResume(err ->
176 | Mono.from(connection.rollbackTransaction()).thenReturn(-1)
177 | ).collect(Collectors.summingInt(i -> (int) i));
178 | }, Connection::close, (c, err) -> c.close(), Connection::close);
179 | }
180 |
181 | }
182 |
--------------------------------------------------------------------------------
/listener-example/src/main/java/io/r2dbc/examples/SpringAopProxyFactory.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import io.r2dbc.proxy.callback.BatchCallbackHandler;
4 | import io.r2dbc.proxy.callback.CallbackHandler;
5 | import io.r2dbc.proxy.callback.ConnectionCallbackHandler;
6 | import io.r2dbc.proxy.callback.ConnectionFactoryCallbackHandler;
7 | import io.r2dbc.proxy.callback.ProxyConfig;
8 | import io.r2dbc.proxy.callback.ResultCallbackHandler;
9 | import io.r2dbc.proxy.callback.StatementCallbackHandler;
10 | import io.r2dbc.proxy.core.ConnectionInfo;
11 | import io.r2dbc.proxy.core.QueryExecutionInfo;
12 | import io.r2dbc.proxy.core.StatementInfo;
13 | import io.r2dbc.spi.Batch;
14 | import io.r2dbc.spi.Connection;
15 | import io.r2dbc.spi.ConnectionFactory;
16 | import io.r2dbc.spi.Result;
17 | import io.r2dbc.spi.Statement;
18 | import io.r2dbc.spi.Wrapped;
19 | import org.aopalliance.intercept.MethodInterceptor;
20 | import org.aopalliance.intercept.MethodInvocation;
21 |
22 | import org.springframework.aop.framework.ProxyFactory;
23 |
24 | /**
25 | * {@link io.r2dbc.proxy.callback.ProxyFactory} implementation that uses spring's {@link ProxyFactory} to create proxy.
26 | *
27 | * @author Tadaya Tsuyukubo
28 | */
29 | public class SpringAopProxyFactory implements io.r2dbc.proxy.callback.ProxyFactory {
30 |
31 | private ProxyConfig proxyConfig;
32 |
33 | public SpringAopProxyFactory(ProxyConfig proxyConfig) {
34 | this.proxyConfig = proxyConfig;
35 | }
36 |
37 | /**
38 | * Interceptor for proxy.
39 | *
40 | * Delegate the invocation to the provided {@link CallbackHandler}.
41 | */
42 | private static class ProxyInterceptor implements MethodInterceptor {
43 | CallbackHandler callbackHandler;
44 |
45 | public ProxyInterceptor(CallbackHandler callbackHandler) {
46 | this.callbackHandler = callbackHandler;
47 | }
48 |
49 | @Override
50 | public Object invoke(MethodInvocation methodInvocation) throws Throwable {
51 | return this.callbackHandler.invoke(methodInvocation.getThis(), methodInvocation.getMethod(), methodInvocation.getArguments());
52 | }
53 | }
54 |
55 | private T createProxy(CallbackHandler callbackHandler, Object target, Class proxyInterface) {
56 | ProxyInterceptor interceptor = new ProxyInterceptor(callbackHandler);
57 |
58 | // NOTE: This ProxyFactory will use jdk dynamic proxy.
59 | // This is because we try to make a proxy on interface, and spring's ProxyFactory
60 | // uses JdkDynamicAopProxy for it.
61 | // See logic detail on "DefaultAopProxyFactory#createAopProxy"
62 | // We could put the actual object and instruct cglib to subclass it; however,
63 | // r2dbc implementations(in this case, H2 driver implementation classes) are
64 | // final classes and cglib cannot subclass final classes.
65 | // Since this implementation is to demonstrate applying different proxy mechanism,
66 | // it is ok to use jdk dynamic proxy.
67 |
68 | ProxyFactory proxyFactory = new ProxyFactory(target);
69 | proxyFactory.addAdvice(interceptor);
70 | proxyFactory.addInterface(proxyInterface);
71 | proxyFactory.addInterface(Wrapped.class); // add this to all proxies
72 | T proxy = proxyInterface.cast(proxyFactory.getProxy());
73 |
74 | return proxy;
75 | }
76 |
77 | @Override
78 | public ConnectionFactory wrapConnectionFactory(ConnectionFactory connectionFactory) {
79 | ConnectionFactoryCallbackHandler handler = new ConnectionFactoryCallbackHandler(connectionFactory, this.proxyConfig);
80 | return createProxy(handler, connectionFactory, ConnectionFactory.class);
81 | }
82 |
83 | @Override
84 | public Connection wrapConnection(Connection connection, ConnectionInfo connectionInfo) {
85 | ConnectionCallbackHandler handler = new ConnectionCallbackHandler(connection, connectionInfo, this.proxyConfig);
86 | return createProxy(handler, connection, Connection.class);
87 | }
88 |
89 | @Override
90 | public Batch wrapBatch(Batch batch, ConnectionInfo connectionInfo) {
91 | BatchCallbackHandler handler = new BatchCallbackHandler(batch, connectionInfo, this.proxyConfig);
92 | return createProxy(handler, batch, Batch.class);
93 | }
94 |
95 | @Override
96 | public Statement wrapStatement(Statement statement, StatementInfo statementInfo, ConnectionInfo connectionInfo) {
97 | StatementCallbackHandler handler = new StatementCallbackHandler(statement, statementInfo, connectionInfo, this.proxyConfig);
98 | return createProxy(handler, statement, Statement.class);
99 | }
100 |
101 | @Override
102 | public Result wrapResult(Result result, QueryExecutionInfo queryExecutionInfo) {
103 | ResultCallbackHandler handler = new ResultCallbackHandler(result, queryExecutionInfo, this.proxyConfig);
104 | return createProxy(handler, result, Result.class);
105 | }
106 |
107 | }
108 |
--------------------------------------------------------------------------------
/listener-example/src/main/java/io/r2dbc/examples/TracingExecutionListener.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import brave.Span;
4 | import brave.Tracer;
5 | import io.r2dbc.proxy.core.*;
6 | import io.r2dbc.proxy.listener.ProxyMethodExecutionListener;
7 |
8 | import static java.util.stream.Collectors.joining;
9 |
10 | /**
11 | * Listener to create spans for R2DBC SPI operations.
12 | *
13 | * @author Tadaya Tsuyukubo
14 | */
15 | public class TracingExecutionListener implements ProxyMethodExecutionListener {
16 |
17 | private static final String TAG_CONNECTION_ID = "connectionId";
18 | private static final String TAG_CONNECTION_CREATE_THREAD_ID = "threadIdOnCreate";
19 | private static final String TAG_CONNECTION_CLOSE_THREAD_ID = "threadIdOnClose";
20 | private static final String TAG_CONNECTION_CREATE_THREAD_NAME = "threadNameOnCreate";
21 | private static final String TAG_CONNECTION_CLOSE_THREAD_NAME = "threadNameOnClose";
22 | private static final String TAG_THREAD_ID = "threadId";
23 | private static final String TAG_THREAD_NAME = "threadName";
24 | private static final String TAG_QUERIES = "queries";
25 | private static final String TAG_BATCH_SIZE = "batchSize";
26 | private static final String TAG_QUERY_TYPE = "type";
27 | private static final String TAG_QUERY_SUCCESS = "success";
28 | private static final String TAG_QUERY_MAPPED_RESULT_COUNT = "mappedResultCount";
29 | private static final String TAG_TRANSACTION_SAVEPOINT = "savepoint";
30 | private static final String TAG_TRANSACTION_COUNT = "transactionCount";
31 | private static final String TAG_COMMIT_COUNT = "commitCount";
32 | private static final String TAG_ROLLBACK_COUNT = "rollbackCount";
33 |
34 | static final String CONNECTION_SPAN_KEY = "connectionSpan";
35 | static final String TRANSACTION_SPAN_KEY = "transactionSpan";
36 | static final String QUERY_SPAN_KEY = "querySpan";
37 |
38 | private final Tracer tracer;
39 |
40 | public TracingExecutionListener(Tracer tracer) {
41 | this.tracer = tracer;
42 | }
43 |
44 | @Override
45 | public void beforeCreateOnConnectionFactory(MethodExecutionInfo methodExecutionInfo) {
46 | Span connectionSpan = this.tracer.nextSpan()
47 | .name("r2dbc:connection")
48 | .kind(Span.Kind.CLIENT)
49 | .start();
50 |
51 | // store the span for retrieval at "afterCreateOnConnectionFactory"
52 | methodExecutionInfo.getValueStore().put("initialConnectionSpan", connectionSpan);
53 | }
54 |
55 | @Override
56 | public void afterCreateOnConnectionFactory(MethodExecutionInfo methodExecutionInfo) {
57 | // retrieve the span created at "beforeCreateOnConnectionFactory"
58 | Span connectionSpan = methodExecutionInfo.getValueStore().get("initialConnectionSpan", Span.class);
59 |
60 | Throwable thrown = methodExecutionInfo.getThrown();
61 | if (thrown != null) {
62 | connectionSpan
63 | .error(thrown)
64 | .finish();
65 | return;
66 | }
67 |
68 | ConnectionInfo connectionInfo = methodExecutionInfo.getConnectionInfo();
69 | String connectionId = connectionInfo.getConnectionId();
70 |
71 | connectionSpan
72 | .tag(TAG_CONNECTION_ID, connectionId)
73 | .tag(TAG_CONNECTION_CREATE_THREAD_ID, String.valueOf(methodExecutionInfo.getThreadId()))
74 | .tag(TAG_CONNECTION_CREATE_THREAD_NAME, methodExecutionInfo.getThreadName())
75 | .annotate("Connection created");
76 |
77 | // store the span in connection scoped value store
78 | connectionInfo.getValueStore().put(CONNECTION_SPAN_KEY, connectionSpan);
79 | }
80 |
81 | @Override
82 | public void afterCloseOnConnection(MethodExecutionInfo methodExecutionInfo) {
83 | ConnectionInfo connectionInfo = methodExecutionInfo.getConnectionInfo();
84 | String connectionId = connectionInfo.getConnectionId();
85 | Span connectionSpan = connectionInfo.getValueStore().get(CONNECTION_SPAN_KEY, Span.class);
86 | if (connectionSpan == null) {
87 | return; // already closed
88 | }
89 | Throwable thrown = methodExecutionInfo.getThrown();
90 | if (thrown != null) {
91 | connectionSpan.error(thrown);
92 | }
93 | connectionSpan
94 | .tag(TAG_CONNECTION_ID, connectionId)
95 | .tag(TAG_CONNECTION_CLOSE_THREAD_ID, String.valueOf(methodExecutionInfo.getThreadId()))
96 | .tag(TAG_CONNECTION_CLOSE_THREAD_NAME, methodExecutionInfo.getThreadName())
97 | .tag(TAG_TRANSACTION_COUNT, String.valueOf(connectionInfo.getTransactionCount()))
98 | .tag(TAG_COMMIT_COUNT, String.valueOf(connectionInfo.getCommitCount()))
99 | .tag(TAG_ROLLBACK_COUNT, String.valueOf(connectionInfo.getRollbackCount()))
100 | .finish();
101 | }
102 |
103 | @Override
104 | public void beforeQuery(QueryExecutionInfo queryExecutionInfo) {
105 | String connectionId = queryExecutionInfo.getConnectionInfo().getConnectionId();
106 |
107 | String queries = queryExecutionInfo.getQueries().stream()
108 | .map(QueryInfo::getQuery)
109 | .collect(joining(", "));
110 |
111 | Span querySpan = this.tracer
112 | .nextSpan()
113 | .name("r2dbc:query")
114 | .kind(Span.Kind.CLIENT)
115 | .tag(TAG_CONNECTION_ID, connectionId)
116 | .tag(TAG_QUERY_TYPE, queryExecutionInfo.getType().toString())
117 | .tag(TAG_QUERIES, queries)
118 | .start();
119 |
120 | if (ExecutionType.BATCH == queryExecutionInfo.getType()) {
121 | querySpan.tag(TAG_BATCH_SIZE, Integer.toString(queryExecutionInfo.getBatchSize()));
122 | }
123 |
124 | // pass the query span to "afterQuery" method
125 | queryExecutionInfo.getValueStore().put(QUERY_SPAN_KEY, querySpan);
126 | }
127 |
128 | @Override
129 | public void afterQuery(QueryExecutionInfo queryExecutionInfo) {
130 | Span querySpan = queryExecutionInfo.getValueStore().get(QUERY_SPAN_KEY, Span.class);
131 | querySpan
132 | .tag(TAG_THREAD_ID, String.valueOf(queryExecutionInfo.getThreadId()))
133 | .tag(TAG_THREAD_NAME, queryExecutionInfo.getThreadName())
134 | .tag(TAG_QUERY_SUCCESS, Boolean.toString(queryExecutionInfo.isSuccess()));
135 |
136 | Throwable thrown = queryExecutionInfo.getThrowable();
137 | if (thrown != null) {
138 | querySpan.error(thrown);
139 | } else {
140 | querySpan.tag(TAG_QUERY_MAPPED_RESULT_COUNT, Integer.toString(queryExecutionInfo.getCurrentResultCount()));
141 | }
142 | querySpan.finish();
143 | }
144 |
145 | @Override
146 | public void beforeBeginTransactionOnConnection(MethodExecutionInfo methodExecutionInfo) {
147 | Span transactionSpan = this.tracer.nextSpan()
148 | .name("r2dbc:transaction")
149 | .kind(Span.Kind.CLIENT)
150 | .start();
151 |
152 | methodExecutionInfo.getConnectionInfo().getValueStore().put(TRANSACTION_SPAN_KEY, transactionSpan);
153 | }
154 |
155 | @Override
156 | public void afterCommitTransactionOnConnection(MethodExecutionInfo methodExecutionInfo) {
157 | ConnectionInfo connectionInfo = methodExecutionInfo.getConnectionInfo();
158 | String connectionId = connectionInfo.getConnectionId();
159 |
160 | Span transactionSpan = connectionInfo.getValueStore().get(TRANSACTION_SPAN_KEY, Span.class);
161 | if (transactionSpan != null) {
162 | transactionSpan
163 | .annotate("Commit")
164 | .tag(TAG_CONNECTION_ID, connectionId)
165 | .tag(TAG_THREAD_ID, String.valueOf(methodExecutionInfo.getThreadId()))
166 | .tag(TAG_THREAD_NAME, methodExecutionInfo.getThreadName())
167 | .finish();
168 | }
169 |
170 | Span connectionSpan = connectionInfo.getValueStore().get(CONNECTION_SPAN_KEY, Span.class);
171 | if (connectionSpan == null) {
172 | return;
173 | }
174 | connectionSpan.annotate("Transaction commit");
175 | }
176 |
177 | @Override
178 | public void afterRollbackTransactionOnConnection(MethodExecutionInfo methodExecutionInfo) {
179 | ConnectionInfo connectionInfo = methodExecutionInfo.getConnectionInfo();
180 | String connectionId = connectionInfo.getConnectionId();
181 |
182 | Span transactionSpan = connectionInfo.getValueStore().get(TRANSACTION_SPAN_KEY, Span.class);
183 | if (transactionSpan != null) {
184 | transactionSpan
185 | .annotate("Rollback")
186 | .tag(TAG_CONNECTION_ID, connectionId)
187 | .tag(TAG_THREAD_ID, String.valueOf(methodExecutionInfo.getThreadId()))
188 | .tag(TAG_THREAD_NAME, methodExecutionInfo.getThreadName())
189 | .finish();
190 | }
191 |
192 | Span connectionSpan = connectionInfo.getValueStore().get(CONNECTION_SPAN_KEY, Span.class);
193 | connectionSpan.annotate("Transaction rollback");
194 | }
195 |
196 | @Override
197 | public void afterRollbackTransactionToSavepointOnConnection(MethodExecutionInfo methodExecutionInfo) {
198 | ConnectionInfo connectionInfo = methodExecutionInfo.getConnectionInfo();
199 | String connectionId = connectionInfo.getConnectionId();
200 | String savepoint = (String) methodExecutionInfo.getMethodArgs()[0];
201 |
202 | Span transactionSpan = connectionInfo.getValueStore().get(TRANSACTION_SPAN_KEY, Span.class);
203 | if (transactionSpan != null) {
204 | transactionSpan
205 | .annotate("Rollback to savepoint")
206 | .tag(TAG_TRANSACTION_SAVEPOINT, savepoint)
207 | .tag(TAG_CONNECTION_ID, connectionId)
208 | .tag(TAG_THREAD_ID, String.valueOf(methodExecutionInfo.getThreadId()))
209 | .tag(TAG_THREAD_NAME, methodExecutionInfo.getThreadName())
210 | .finish();
211 | }
212 |
213 | Span connectionSpan = connectionInfo.getValueStore().get(CONNECTION_SPAN_KEY, Span.class);
214 | connectionSpan.annotate("Transaction rollback to savepoint");
215 | }
216 |
217 | }
218 |
--------------------------------------------------------------------------------
/listener-example/src/main/resources/application.yaml:
--------------------------------------------------------------------------------
1 | spring:
2 | application:
3 | name: r2dbc-proxy-examples
4 | sleuth:
5 | sampler:
6 | probability: 1.0
7 |
8 | management:
9 | endpoint:
10 | metrics:
11 | enabled: true
12 | endpoints:
13 | web:
14 | exposure:
15 | include: metrics
16 |
--------------------------------------------------------------------------------
/listener-example/src/test/java/io/r2dbc/examples/MetricsExecutionListenerTest.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import io.micrometer.core.instrument.Counter;
4 | import io.micrometer.core.instrument.Meter;
5 | import io.micrometer.core.instrument.Tag;
6 | import io.micrometer.core.instrument.Timer;
7 | import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
8 | import io.r2dbc.proxy.core.DefaultValueStore;
9 | import io.r2dbc.proxy.core.ValueStore;
10 | import io.r2dbc.proxy.test.MockMethodExecutionInfo;
11 | import io.r2dbc.proxy.test.MockQueryExecutionInfo;
12 | import org.junit.jupiter.api.BeforeEach;
13 | import org.junit.jupiter.api.Test;
14 |
15 | import java.util.List;
16 |
17 | import static org.assertj.core.api.Assertions.assertThat;
18 |
19 | /**
20 | * Test for {@link MetricsExecutionListener}.
21 | *
22 | * @author Tadaya Tsuyukubo
23 | */
24 | class MetricsExecutionListenerTest {
25 |
26 | private SimpleMeterRegistry registry;
27 | private MetricsExecutionListener listener;
28 |
29 | @BeforeEach
30 | void beforeEach() {
31 | this.registry = new SimpleMeterRegistry();
32 | this.listener = new MetricsExecutionListener(this.registry);
33 | }
34 |
35 | @Test
36 | void createConnection() {
37 | ValueStore valueStore = new DefaultValueStore();
38 | MockMethodExecutionInfo executionInfo = MockMethodExecutionInfo.builder().valueStore(valueStore).build();
39 |
40 | this.listener.beforeCreateOnConnectionFactory(executionInfo);
41 | this.listener.afterCreateOnConnectionFactory(executionInfo);
42 |
43 | List meters = this.registry.getMeters();
44 | assertThat(meters).hasSize(1)
45 | .first()
46 | .isInstanceOfSatisfying(Timer.class, (timer) -> {
47 | assertThat(timer.getId().getName()).isEqualTo("r2dbc.connection");
48 | assertThat(timer.getId().getTags()).containsExactly(Tag.of("event", "create"));
49 | });
50 | }
51 |
52 | @Test
53 | void commit() {
54 | MockMethodExecutionInfo executionInfo = MockMethodExecutionInfo.empty();
55 | this.listener.afterCommitTransactionOnConnection(executionInfo);
56 |
57 | List meters = this.registry.getMeters();
58 | assertThat(meters).hasSize(1)
59 | .first()
60 | .isInstanceOfSatisfying(Counter.class, (counter) -> {
61 | assertThat(counter.getId().getName()).isEqualTo("r2dbc.transaction");
62 | assertThat(counter.getId().getTags()).containsExactly(Tag.of("event", "commit"));
63 | assertThat(counter.count()).isEqualTo(1);
64 | });
65 | }
66 |
67 | @Test
68 | void rollback() {
69 | MockMethodExecutionInfo executionInfo = MockMethodExecutionInfo.empty();
70 | this.listener.afterRollbackTransactionOnConnection(executionInfo);
71 |
72 | List meters = this.registry.getMeters();
73 | assertThat(meters).hasSize(1)
74 | .first()
75 | .isInstanceOfSatisfying(Counter.class, (counter) -> {
76 | assertThat(counter.getId().getName()).isEqualTo("r2dbc.transaction");
77 | assertThat(counter.getId().getTags()).containsExactly(Tag.of("event", "rollback"));
78 | assertThat(counter.count()).isEqualTo(1);
79 | });
80 | }
81 |
82 | @Test
83 | void afterExecuteOnBatch() {
84 | MockQueryExecutionInfo executionInfo = MockQueryExecutionInfo.empty();
85 | this.listener.afterExecuteOnBatch(executionInfo);
86 |
87 | List meters = this.registry.getMeters();
88 | assertThat(meters).hasSize(1)
89 | .first()
90 | .isInstanceOfSatisfying(Counter.class, (counter) -> {
91 | assertThat(counter.getId().getName()).isEqualTo("r2dbc.query");
92 | assertThat(counter.count()).isEqualTo(1);
93 | });
94 | }
95 |
96 | @Test
97 | void afterExecuteOnStatement() {
98 | MockQueryExecutionInfo executionInfo = MockQueryExecutionInfo.empty();
99 | this.listener.afterExecuteOnStatement(executionInfo);
100 |
101 | List meters = this.registry.getMeters();
102 | assertThat(meters).hasSize(1)
103 | .first()
104 | .isInstanceOfSatisfying(Counter.class, (counter) -> {
105 | assertThat(counter.getId().getName()).isEqualTo("r2dbc.query");
106 | assertThat(counter.count()).isEqualTo(1);
107 | });
108 | }
109 |
110 | }
111 |
--------------------------------------------------------------------------------
/listener-example/src/test/java/io/r2dbc/examples/QueryTimeMetricsExecutionListenerTest.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import io.micrometer.core.instrument.Timer;
4 | import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
5 | import io.r2dbc.proxy.core.QueryExecutionInfo;
6 | import io.r2dbc.proxy.core.QueryInfo;
7 | import io.r2dbc.proxy.test.MockQueryExecutionInfo;
8 | import org.junit.jupiter.api.extension.ExtensionContext;
9 | import org.junit.jupiter.params.ParameterizedTest;
10 | import org.junit.jupiter.params.provider.Arguments;
11 | import org.junit.jupiter.params.provider.ArgumentsProvider;
12 | import org.junit.jupiter.params.provider.ArgumentsSource;
13 |
14 | import java.time.Duration;
15 | import java.util.concurrent.TimeUnit;
16 | import java.util.stream.Stream;
17 |
18 | import static org.assertj.core.api.Assertions.assertThat;
19 | import static org.junit.jupiter.params.provider.Arguments.arguments;
20 |
21 | /**
22 | * Test for {@link QueryTimeMetricsExecutionListener}.
23 | *
24 | * @author Tadaya Tsuyukubo
25 | */
26 | class QueryTimeMetricsExecutionListenerTest {
27 |
28 | @ParameterizedTest
29 | @ArgumentsSource(MetricsArgumentsProvider.class)
30 | void metrics(String query, String meterName) {
31 | SimpleMeterRegistry registry = new SimpleMeterRegistry();
32 | QueryTimeMetricsExecutionListener listener = new QueryTimeMetricsExecutionListener(registry);
33 |
34 | QueryInfo queryInfo = new QueryInfo(query);
35 | QueryExecutionInfo queryExecutionInfo = new MockQueryExecutionInfo.Builder()
36 | .queryInfo(queryInfo)
37 | .executeDuration(Duration.ofSeconds(10))
38 | .build();
39 |
40 | listener.afterQuery(queryExecutionInfo);
41 |
42 | Timer select = registry.get(meterName).timer();
43 | assertThat(select.count()).isEqualTo(1);
44 | assertThat(select.totalTime(TimeUnit.SECONDS)).isEqualTo(10);
45 | }
46 |
47 | private static class MetricsArgumentsProvider implements ArgumentsProvider {
48 |
49 | @Override
50 | public Stream provideArguments(ExtensionContext context) {
51 | return Stream.of(
52 | // query, meter name
53 | arguments("SELECT ...", "r2dbc.query.select"),
54 | arguments("INSERT ...", "r2dbc.query.insert"),
55 | arguments("UPDATE ...", "r2dbc.query.update"),
56 | arguments("DELETE ...", "r2dbc.query.delete"),
57 | arguments("UPSERT ...", "r2dbc.query.other")
58 | );
59 | }
60 | }
61 | }
62 |
--------------------------------------------------------------------------------
/listener-example/src/test/java/io/r2dbc/examples/TracingExecutionListenerTest.java:
--------------------------------------------------------------------------------
1 | package io.r2dbc.examples;
2 |
3 | import brave.Span;
4 | import brave.Tracer;
5 | import brave.Tracing;
6 | import brave.propagation.StrictCurrentTraceContext;
7 | import brave.test.TestSpanHandler;
8 | import io.r2dbc.proxy.core.ConnectionInfo;
9 | import io.r2dbc.proxy.core.ExecutionType;
10 | import io.r2dbc.proxy.core.QueryInfo;
11 | import io.r2dbc.proxy.core.ValueStore;
12 | import io.r2dbc.proxy.test.MockConnectionInfo;
13 | import io.r2dbc.proxy.test.MockMethodExecutionInfo;
14 | import io.r2dbc.proxy.test.MockQueryExecutionInfo;
15 | import org.junit.jupiter.api.AfterEach;
16 | import org.junit.jupiter.api.Test;
17 |
18 | import static io.r2dbc.examples.TracingExecutionListener.CONNECTION_SPAN_KEY;
19 | import static io.r2dbc.examples.TracingExecutionListener.TRANSACTION_SPAN_KEY;
20 | import static org.assertj.core.api.Assertions.assertThat;
21 |
22 | /**
23 | * Test for {@link TracingExecutionListener}.
24 | *
25 | * @author Tadaya Tsuyukubo
26 | */
27 | class TracingExecutionListenerTest {
28 |
29 | private StrictCurrentTraceContext currentTraceContext = StrictCurrentTraceContext.create();
30 | private TestSpanHandler spanHandler = new TestSpanHandler();
31 | private Tracing tracing = Tracing.newBuilder()
32 | .currentTraceContext(currentTraceContext)
33 | .addSpanHandler(spanHandler)
34 | .build();
35 |
36 | private Tracer tracer = tracing.tracer();
37 | private TracingExecutionListener listener = new TracingExecutionListener(tracer);
38 |
39 | @AfterEach
40 | void afterEach() {
41 | Tracing.current().close();
42 | }
43 |
44 | @Test
45 | void query() {
46 | ValueStore valueStore = ValueStore.create();
47 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
48 | .connectionId("foo")
49 | .valueStore(valueStore)
50 | .build();
51 | QueryInfo queryInfo = new QueryInfo("SELECT 1");
52 | MockQueryExecutionInfo queryExecutionInfo = MockQueryExecutionInfo.builder()
53 | .connectionInfo(connectionInfo)
54 | .queryInfo(queryInfo)
55 | .type(ExecutionType.STATEMENT)
56 | .threadName("thread-name")
57 | .threadId(300)
58 | .isSuccess(true)
59 | .build();
60 |
61 | this.listener.beforeQuery(queryExecutionInfo);
62 | this.listener.afterQuery(queryExecutionInfo);
63 |
64 | assertThat(this.spanHandler.spans()).hasSize(1);
65 | assertThat(this.spanHandler.get(0).name()).isEqualTo("r2dbc:query");
66 | assertThat(this.spanHandler.get(0).tags())
67 | .containsEntry("queries", "SELECT 1")
68 | .containsEntry("threadName", "thread-name")
69 | .containsEntry("threadId", "300")
70 | .containsEntry("success", "true")
71 | ;
72 | }
73 |
74 | @Test
75 | void createConnection() {
76 | ValueStore valueStore = ValueStore.create();
77 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
78 | .connectionId("foo")
79 | .valueStore(valueStore)
80 | .build();
81 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
82 | .connectionInfo(connectionInfo)
83 | .threadId(10)
84 | .threadName("thread-name")
85 | .build();
86 |
87 | this.listener.beforeCreateOnConnectionFactory(methodExecutionInfo);
88 | this.listener.afterCreateOnConnectionFactory(methodExecutionInfo);
89 |
90 |
91 | assertThat(valueStore.get(CONNECTION_SPAN_KEY)).as("Connection span should be stored")
92 | .isNotNull().isInstanceOf(Span.class);
93 |
94 | Span span = valueStore.get(CONNECTION_SPAN_KEY, Span.class);
95 | span.finish();
96 |
97 | assertThat(this.spanHandler.spans()).hasSize(1);
98 | assertThat(this.spanHandler.get(0).name()).isEqualTo("r2dbc:connection");
99 | assertThat(this.spanHandler.get(0).tags())
100 | .containsEntry("connectionId", "foo")
101 | .containsEntry("threadNameOnCreate", "thread-name")
102 | .containsEntry("threadIdOnCreate", "10")
103 | ;
104 | assertThat(this.spanHandler.get(0).containsAnnotation("Connection created")).isTrue();
105 | }
106 |
107 | @Test
108 | void createConnectionWithError() {
109 | Exception error = new RuntimeException();
110 |
111 | ValueStore valueStore = ValueStore.create();
112 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
113 | .connectionId("foo")
114 | .valueStore(valueStore)
115 | .build();
116 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
117 | .connectionInfo(connectionInfo)
118 | .threadId(10)
119 | .threadName("thread-name")
120 | .setThrown(error)
121 | .build();
122 |
123 | this.listener.beforeCreateOnConnectionFactory(methodExecutionInfo);
124 | this.listener.afterCreateOnConnectionFactory(methodExecutionInfo);
125 |
126 | assertThat(this.spanHandler.spans()).hasSize(1);
127 | assertThat(this.spanHandler.get(0).name()).isEqualTo("r2dbc:connection");
128 | assertThat(this.spanHandler.get(0).error()).isSameAs(error);
129 | }
130 |
131 | @Test
132 | void closeConnection() {
133 | Span span = this.tracer.nextSpan().kind(Span.Kind.CLIENT).start();
134 |
135 | ValueStore valueStore = ValueStore.create();
136 | valueStore.put(CONNECTION_SPAN_KEY, span);
137 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
138 | .connectionId("foo")
139 | .commitCount(10)
140 | .rollbackCount(20)
141 | .transactionCount(30)
142 | .valueStore(valueStore)
143 | .build();
144 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
145 | .connectionInfo(connectionInfo)
146 | .threadName("thread-name")
147 | .threadId(300)
148 | .build();
149 |
150 | this.listener.afterCloseOnConnection(methodExecutionInfo);
151 |
152 | assertThat(this.spanHandler.spans()).hasSize(1);
153 | assertThat(this.spanHandler.get(0).tags())
154 | .containsEntry("connectionId", "foo")
155 | .containsEntry("threadNameOnClose", "thread-name")
156 | .containsEntry("threadIdOnClose", "300")
157 | .containsEntry("commitCount", "10")
158 | .containsEntry("rollbackCount", "20")
159 | .containsEntry("transactionCount", "30")
160 | ;
161 | }
162 |
163 | @Test
164 | void closeConnectionWithError() {
165 | Span span = this.tracer.nextSpan().kind(Span.Kind.CLIENT).start();
166 |
167 | Exception error = new RuntimeException();
168 |
169 | ValueStore valueStore = ValueStore.create();
170 | valueStore.put(CONNECTION_SPAN_KEY, span);
171 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
172 | .connectionId("foo")
173 | .valueStore(valueStore)
174 | .build();
175 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
176 | .connectionInfo(connectionInfo)
177 | .threadName("thread-name")
178 | .setThrown(error)
179 | .build();
180 |
181 | this.listener.afterCloseOnConnection(methodExecutionInfo);
182 |
183 | assertThat(this.spanHandler.spans()).hasSize(1);
184 | assertThat(this.spanHandler.get(0).error()).isSameAs(error);
185 | }
186 |
187 | @Test
188 | void beginTransaction() {
189 | ValueStore valueStore = ValueStore.create();
190 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
191 | .connectionId("foo")
192 | .valueStore(valueStore)
193 | .build();
194 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
195 | .connectionInfo(connectionInfo)
196 | .build();
197 |
198 | this.listener.beforeBeginTransactionOnConnection(methodExecutionInfo);
199 |
200 | Span span = valueStore.get(TRANSACTION_SPAN_KEY, Span.class);
201 | assertThat(span).as("transaction span should be stored").isNotNull();
202 |
203 | assertThat(this.spanHandler.spans()).as("Span is not finished yet").isEmpty();
204 |
205 | span.finish();
206 | assertThat(this.spanHandler.get(0).name()).isEqualTo("r2dbc:transaction");
207 | }
208 |
209 | @Test
210 | void transactionCommit() {
211 | Span connSpan = this.tracer.nextSpan().kind(Span.Kind.CLIENT).start();
212 | Span txSpan = this.tracer.nextSpan().kind(Span.Kind.CLIENT).start();
213 |
214 | ValueStore valueStore = ValueStore.create();
215 | valueStore.put(CONNECTION_SPAN_KEY, connSpan);
216 | valueStore.put(TRANSACTION_SPAN_KEY, txSpan);
217 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
218 | .connectionId("foo")
219 | .valueStore(valueStore)
220 | .build();
221 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
222 | .connectionInfo(connectionInfo)
223 | .threadId(10)
224 | .threadName("thread-name")
225 | .build();
226 |
227 | this.listener.afterCommitTransactionOnConnection(methodExecutionInfo);
228 |
229 | // check txSpan
230 | assertThat(this.spanHandler.spans()).hasSize(1);
231 | assertThat(this.spanHandler.get(0).tags())
232 | .containsEntry("connectionId", "foo")
233 | .containsEntry("threadName", "thread-name")
234 | .containsEntry("threadId", "10")
235 | ;
236 | assertThat(this.spanHandler.get(0).containsAnnotation("Commit")).isTrue();
237 |
238 | // check connSpan
239 | this.spanHandler.clear();
240 | connSpan.finish();
241 | assertThat(this.spanHandler.spans()).hasSize(1);
242 | assertThat(this.spanHandler.get(0).containsAnnotation("Transaction commit")).isTrue();
243 | }
244 |
245 | @Test
246 | void transactionRollback() {
247 | Span connSpan = this.tracer.nextSpan().kind(Span.Kind.CLIENT).start();
248 | Span txSpan = this.tracer.nextSpan().kind(Span.Kind.CLIENT).start();
249 |
250 | ValueStore valueStore = ValueStore.create();
251 | valueStore.put(CONNECTION_SPAN_KEY, connSpan);
252 | valueStore.put(TRANSACTION_SPAN_KEY, txSpan);
253 | ConnectionInfo connectionInfo = MockConnectionInfo.builder()
254 | .connectionId("foo")
255 | .valueStore(valueStore)
256 | .build();
257 | MockMethodExecutionInfo methodExecutionInfo = MockMethodExecutionInfo.builder()
258 | .connectionInfo(connectionInfo)
259 | .threadId(10)
260 | .threadName("thread-name")
261 | .build();
262 |
263 | this.listener.afterRollbackTransactionOnConnection(methodExecutionInfo);
264 |
265 | // check txSpan
266 | assertThat(this.spanHandler.spans()).hasSize(1);
267 | assertThat(this.spanHandler.get(0).tags())
268 | .containsEntry("connectionId", "foo")
269 | .containsEntry("threadName", "thread-name")
270 | .containsEntry("threadId", "10")
271 | ;
272 | assertThat(this.spanHandler.get(0).containsAnnotation("Rollback")).isTrue();
273 |
274 | // check connSpan
275 | this.spanHandler.clear();
276 | connSpan.finish();
277 | assertThat(this.spanHandler.spans()).hasSize(1);
278 | assertThat(this.spanHandler.get(0).containsAnnotation("Transaction rollback")).isTrue();
279 | }
280 |
281 | }
282 |
--------------------------------------------------------------------------------