Use beforeQuery/afterQuery methods for query execution callback

ttddyy · Jul 2, 2020 · 16d30f6 · 16d30f6
1 parent 9a9c21c
commit 16d30f6
Showing 1 changed file with 8 additions and 30 deletions.
diff --git a/listener-example/src/main/java/io/r2dbc/examples/TracingExecutionListener.java b/listener-example/src/main/java/io/r2dbc/examples/TracingExecutionListener.java
@@ -1,17 +1,13 @@
 package io.r2dbc.examples;
 
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
 import brave.Span;
 import brave.Tracer;
-import io.r2dbc.proxy.core.ConnectionInfo;
-import io.r2dbc.proxy.core.ExecutionType;
-import io.r2dbc.proxy.core.MethodExecutionInfo;
-import io.r2dbc.proxy.core.QueryExecutionInfo;
-import io.r2dbc.proxy.core.QueryInfo;
+import io.r2dbc.proxy.core.*;
 import io.r2dbc.proxy.listener.LifeCycleListener;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
 import static java.util.stream.Collectors.joining;
 
 /**
@@ -100,16 +96,7 @@ public void afterCloseOnConnection(MethodExecutionInfo methodExecutionInfo) {
     }
 
     @Override
-    public void beforeExecuteOnBatch(QueryExecutionInfo queryExecutionInfo) {
-        beforeExecuteQuery(queryExecutionInfo);
-    }
-
-    @Override
-    public void beforeExecuteOnStatement(QueryExecutionInfo queryExecutionInfo) {
-        beforeExecuteQuery(queryExecutionInfo);
-    }
-
-    private void beforeExecuteQuery(QueryExecutionInfo queryExecutionInfo) {
+    public void beforeQuery(QueryExecutionInfo queryExecutionInfo) {
         String connectionId = queryExecutionInfo.getConnectionInfo().getConnectionId();
 
         String queries = queryExecutionInfo.getQueries().stream()
@@ -129,21 +116,12 @@ private void beforeExecuteQuery(QueryExecutionInfo queryExecutionInfo) {
             querySpan.tag(TAG_BATCH_SIZE, Integer.toString(queryExecutionInfo.getBatchSize()));
         }
 
-		queryExecutionInfo.getValueStore().put("querySpan", querySpan);
-    }
-
-    @Override
-    public void afterExecuteOnBatch(QueryExecutionInfo queryExecutionInfo) {
-        afterExecuteQuery(queryExecutionInfo);
+        queryExecutionInfo.getValueStore().put("querySpan", querySpan);
     }
 
     @Override
-    public void afterExecuteOnStatement(QueryExecutionInfo queryExecutionInfo) {
-        afterExecuteQuery(queryExecutionInfo);
-    }
-
-    private void afterExecuteQuery(QueryExecutionInfo queryExecutionInfo) {
-		Span querySpan = queryExecutionInfo.getValueStore().get("querySpan", Span.class);
+    public void afterQuery(QueryExecutionInfo queryExecutionInfo) {
+        Span querySpan = queryExecutionInfo.getValueStore().get("querySpan", Span.class);
 
         querySpan
                 .tag(TAG_THREAD_ID, String.valueOf(queryExecutionInfo.getThreadId()))