Add proxy object to the method callback parameter

Resolve #108
jdbc-observations · Dec 5, 2023 · 1a42780 · 1a42780
1 parent e7dddfe
commit 1a42780
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 1 deletion.
diff --git a/src/main/java/net/ttddyy/dsproxy/listener/MethodExecutionContext.java b/src/main/java/net/ttddyy/dsproxy/listener/MethodExecutionContext.java
@@ -2,6 +2,7 @@
 
 import net.ttddyy.dsproxy.ConnectionInfo;
 import net.ttddyy.dsproxy.proxy.ProxyConfig;
+import net.ttddyy.dsproxy.proxy.ProxyJdbcObject;
 
 import java.lang.reflect.Method;
 import java.util.HashMap;
@@ -25,6 +26,7 @@ public static class Builder {
         private long elapsedTime;
         private ConnectionInfo connectionInfo;
         private ProxyConfig proxyConfig;
+        private Object proxy;
 
         public static Builder create() {
             return new Builder();
@@ -40,6 +42,7 @@ public MethodExecutionContext build() {
             context.elapsedTime = this.elapsedTime;
             context.connectionInfo = this.connectionInfo;
             context.proxyConfig = this.proxyConfig;
+            context.proxy = this.proxy;
             return context;
         }
 
@@ -82,6 +85,11 @@ public Builder proxyConfig(ProxyConfig proxyConfig) {
             this.proxyConfig = proxyConfig;
             return this;
         }
+
+        public Builder proxy(Object proxy) {
+            this.proxy = proxy;
+            return this;
+        }
     }
 
 
@@ -93,6 +101,7 @@ public Builder proxyConfig(ProxyConfig proxyConfig) {
     private long elapsedTime;
     private ConnectionInfo connectionInfo;
     private ProxyConfig proxyConfig;
+    private Object proxy;
     private Map<String, Object> customValues = new HashMap<String, Object>();
 
     public Object getTarget() {
@@ -190,6 +199,20 @@ public void setProxyConfig(ProxyConfig proxyConfig) {
         this.proxyConfig = proxyConfig;
     }
 
+    /**
+     * Get the proxy object.
+     *
+     * @return proxy object
+     * @since 1.10
+     */
+    public Object getProxy() {
+        return this.proxy;
+    }
+
+    public void setProxy(Object proxy) {
+        this.proxy = proxy;
+    }
+
     /**
      * Store key/value pair.
      *

diff --git a/src/main/java/net/ttddyy/dsproxy/proxy/ProxyLogicSupport.java b/src/main/java/net/ttddyy/dsproxy/proxy/ProxyLogicSupport.java
@@ -82,6 +82,7 @@ protected Object proceedMethodExecution(ProxyConfig proxyConfig, Object original
                                             Object proxy, Method method, Object[] args) throws Throwable {
         MethodExecutionContext methodContext = MethodExecutionContext.Builder.create()
                 .target(original)
+                .proxy(proxy)
                 .method(method)
                 .methodArgs(args)
                 .connectionInfo(connectionInfo)

diff --git a/src/test/java/net/ttddyy/dsproxy/proxy/ProxyLogicSupportTest.java b/src/test/java/net/ttddyy/dsproxy/proxy/ProxyLogicSupportTest.java
@@ -21,6 +21,7 @@ public class ProxyLogicSupportTest {
     @Test
     public void proceedMethodExecution() throws Throwable {
         final Object target = new Object();
+        final Object proxy = new Object();
         final Method method = Statement.class.getMethod("getConnection");
         final Object[] methodArgs = new Object[]{};
         final Object returnObj = new Object();
@@ -41,6 +42,7 @@ public void beforeMethod(MethodExecutionContext executionContext) {
 
                 assertThat(executionContext.getConnectionInfo()).isSameAs(connectionInfo);
                 assertThat(executionContext.getProxyConfig()).isNotNull();
+                assertThat(executionContext.getProxy()).isSameAs(proxy);
             }
 
             @Override
@@ -57,14 +59,15 @@ public void afterMethod(MethodExecutionContext executionContext) {
 
                 assertThat(executionContext.getConnectionInfo()).isSameAs(connectionInfo);
                 assertThat(executionContext.getProxyConfig()).isNotNull();
+                assertThat(executionContext.getProxy()).isSameAs(proxy);
             }
         };
 
 
         ProxyConfig proxyConfig = ProxyConfig.Builder.create().methodListener(listener).build();
 
         Custom custom = new Custom((proxyTarget, m, args) -> returnObj);
-        Object result = custom.proceedMethodExecution(proxyConfig, target, connectionInfo, null, method, methodArgs);
+        Object result = custom.proceedMethodExecution(proxyConfig, target, connectionInfo, proxy, method, methodArgs);
 
         assertSame(returnObj, result);
         assertTrue(listener.isBeforeMethodCalled());