improve java.sql.Wrapper support

src/main/java/com/alibaba/druid/filter/FilterChainImpl.java

@@ -105,7 +105,7 @@ public class FilterChainImpl implements FilterChain {
         }
 
         // // if driver is for jdbc 3.0
-        if (iface.isAssignableFrom(wrapper.getClass())) {
+        if (iface.isInstance(wrapper)) {
             return true;
         }
 

src/main/java/com/alibaba/druid/mock/MockConnection.java

@@ -98,7 +98,7 @@ public class MockConnection extends ConnectionBase implements Connection {
     @SuppressWarnings("unchecked")
     @Override
     public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == MockConnection.class) {
+        if (iface.isInstance(this)) {
             return (T) this;
         }
 
@@ -107,7 +107,7 @@ public class MockConnection extends ConnectionBase implements Connection {
 
     @Override
     public boolean isWrapperFor(Class<?> iface) throws SQLException {
-        return iface == MockConnection.class;
+        return iface.isInstance(this);
     }
 
     @Override

src/main/java/com/alibaba/druid/mock/MockDriver.java

@@ -234,6 +234,10 @@ public class MockDriver implements Driver, MockDriverMBean {
     public boolean jdbcCompliant() {
         return true;
     }
+    
+    public MockResultSet createMockResultSet(MockStatementBase stmt) {
+        return new MockResultSet(stmt);
+    }
 
     public ResultSet executeQuery(MockStatementBase stmt, String sql) throws SQLException {
         if (logExecuteQueryEnable && LOG.isDebugEnabled()) {
@@ -252,7 +256,7 @@ public class MockDriver implements Driver, MockDriverMBean {
         handleSleep(conn);
 
         if ("SELECT value FROM _int_1000_".equalsIgnoreCase(sql)) {
-            MockResultSet rs = new MockResultSet(stmt);
+            MockResultSet rs = createMockResultSet(stmt);
 
             for (int i = 0; i < 1000; ++i) {
                 rs.getRows().add(new Object[] { i });

src/main/java/com/alibaba/druid/mock/MockPreparedStatement.java

@@ -50,6 +50,8 @@ public class MockPreparedStatement extends PreparedStatementBase implements Mock
 
         if (conn != null) {
             conn.handleSleep();
+            
+            return conn.getDriver().createMockResultSet(this);
         }
         
         return new MockResultSet(this);

src/main/java/com/alibaba/druid/mock/handler/MySqlMockExecuteHandlerImpl.java

@@ -120,7 +120,7 @@ public class MySqlMockExecuteHandlerImpl implements MockExecuteHandler {
     }
 
     public ResultSet executeQueryFromDual(MockStatementBase statement, SQLSelectQueryBlock query) throws SQLException {
-        MockResultSet rs = new MockResultSet(statement);
+        MockResultSet rs = statement.getConnection().getDriver().createMockResultSet(statement);
         MockResultSetMetaData metaData = rs.getMockMetaData();
 
         Object[] row = new Object[query.getSelectList().size()];

src/main/java/com/alibaba/druid/pool/DruidPooledConnection.java

@@ -44,7 +44,6 @@ import com.alibaba.druid.filter.Filter;
 import com.alibaba.druid.filter.FilterChainImpl;
 import com.alibaba.druid.pool.DruidPooledPreparedStatement.PreparedStatementKey;
 import com.alibaba.druid.pool.PreparedStatementPool.MethodType;
-import com.alibaba.druid.proxy.jdbc.ConnectionProxy;
 import com.alibaba.druid.proxy.jdbc.TransactionInfo;
 import com.alibaba.druid.support.logging.Log;
 import com.alibaba.druid.support.logging.LogFactory;
@@ -52,7 +51,7 @@ import com.alibaba.druid.support.logging.LogFactory;
 /**
  * @author wenshao<szujobs@hotmail.com>
  */
-public class DruidPooledConnection implements javax.sql.PooledConnection, Connection {
+public class DruidPooledConnection extends PoolableWrapper implements javax.sql.PooledConnection, Connection {
 
     private final static Log                 LOG         = LogFactory.getLog(DruidPooledConnection.class);
 
@@ -73,6 +72,8 @@ public class DruidPooledConnection implements javax.sql.PooledConnection, Connec
     private StackTraceElement[]              connectStackTrace;
 
     public DruidPooledConnection(DruidConnectionHolder holder){
+        super(holder.getConnection());
+        
         this.conn = holder.getConnection();
         this.holder = holder;
         dupCloseLogEnable = holder.getDataSource().isDupCloseLogEnable();
@@ -645,47 +646,6 @@ public class DruidPooledConnection implements javax.sql.PooledConnection, Connec
         return poolableStatement;
     }
 
-    // ////////////////////
-
-    @SuppressWarnings("unchecked")
-    @Override
-    public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == null) {
-            return null;
-        }
-
-        if (iface == Connection.class) {
-            if (conn instanceof ConnectionProxy) {
-                return conn.unwrap(iface);
-            }
-
-            return (T) conn;
-        }
-
-        if (iface.isInstance(conn)) {
-            return (T) conn;
-        }
-
-        if (iface.isInstance(this)) {
-            return (T) this;
-        }
-
-        return conn.unwrap(iface);
-    }
-
-    @Override
-    public boolean isWrapperFor(Class<?> iface) throws SQLException {
-        if (iface == null) {
-            return false;
-        }
-
-        if (iface.isInstance(this)) {
-            return true;
-        }
-
-        return conn.isWrapperFor(iface);
-    }
-
     @Override
     public String nativeSQL(String sql) throws SQLException {
         checkState();

src/main/java/com/alibaba/druid/pool/DruidPooledPreparedStatement.java

@@ -38,7 +38,6 @@ import java.util.Arrays;
 import java.util.Calendar;
 
 import com.alibaba.druid.pool.PreparedStatementPool.MethodType;
-import com.alibaba.druid.proxy.jdbc.PreparedStatementProxy;
 import com.alibaba.druid.support.logging.Log;
 import com.alibaba.druid.support.logging.LogFactory;
 import com.alibaba.druid.util.OracleUtils;
@@ -1035,20 +1034,21 @@ public class DruidPooledPreparedStatement extends DruidPooledStatement implement
         }
 
     }
-
+    
+    @Override
+    public boolean isWrapperFor(Class<?> iface) throws SQLException {
+        if (iface == PreparedStatementHolder.class) {
+            return true;
+        }
+        
+        return super.isWrapperFor(iface);
+    }
+    
     @SuppressWarnings("unchecked")
     public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == PreparedStatement.class) {
-            if (stmt instanceof PreparedStatementProxy) {
-                return stmt.unwrap(iface);
-            }
-            return (T) stmt;
-        }
-
         if (iface == PreparedStatementHolder.class) {
-            return (T) this.holder;
+            return (T) holder;
         }
-
         return super.unwrap(iface);
     }
 }

src/main/java/com/alibaba/druid/pool/DruidPooledResultSet.java

@@ -33,8 +33,6 @@ import java.sql.SQLXML;
 import java.sql.Statement;
 import java.util.Calendar;
 
-import com.alibaba.druid.proxy.jdbc.ResultSetProxy;
-
 /**
  * @author wenshao<szujobs@hotmail.com>
  */
@@ -1772,17 +1770,4 @@ public final class DruidPooledResultSet extends PoolableWrapper implements Resul
     public <T> T getObject(String columnLabel, Class<T> type) throws SQLException {
         throw new SQLFeatureNotSupportedException();
     }
-    
-    @SuppressWarnings("unchecked")
-    public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == ResultSet.class) {
-            if (rs instanceof ResultSetProxy) {
-                return rs.unwrap(iface);
-            }
-            
-            return (T) rs;
-        }
-        
-        return super.unwrap(iface);
-    }
 }

src/main/java/com/alibaba/druid/pool/DruidPooledStatement.java

@@ -24,7 +24,6 @@ import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
 
-import com.alibaba.druid.proxy.jdbc.StatementProxy;
 import com.alibaba.druid.support.logging.Log;
 import com.alibaba.druid.support.logging.LogFactory;
 
@@ -628,16 +627,4 @@ public class DruidPooledStatement extends PoolableWrapper implements Statement {
     public boolean isCloseOnCompletion() throws SQLException {
         throw new SQLFeatureNotSupportedException();
     }
-
-    @SuppressWarnings("unchecked")
-    public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == Statement.class) {
-            if (stmt instanceof StatementProxy) {
-                return stmt.unwrap(iface);
-            }
-            return (T) stmt;
-        }
-
-        return super.unwrap(iface);
-    }
 }

src/main/java/com/alibaba/druid/pool/PoolableWrapper.java

@@ -18,6 +18,8 @@ package com.alibaba.druid.pool;
 import java.sql.SQLException;
 import java.sql.Wrapper;
 
+import com.alibaba.druid.proxy.jdbc.WrapperProxy;
+
 /**
  * @author wenshao<szujobs@hotmail.com>
  */
@@ -35,13 +37,19 @@ public class PoolableWrapper implements Wrapper {
             return false;
         }
 
-        if (iface.isInstance(wrapper)) {
+        if (wrapper != null && iface == wrapper.getClass()) {
             return true;
         }
 
-        if (iface.isInstance(this)) {
+        if (iface == this.getClass()) {
             return true;
         }
+        
+        if (!(wrapper instanceof WrapperProxy)) {
+            if (iface.isInstance(wrapper)) {
+                return true;
+            }
+        }
 
         return wrapper.isWrapperFor(iface);
     }
@@ -53,13 +61,20 @@ public class PoolableWrapper implements Wrapper {
             return null;
         }
 
-        if (iface.isInstance(wrapper)) {
+        if (wrapper != null && iface == wrapper.getClass()) {
             return (T) wrapper;
         }
 
-        if (iface.isInstance(this)) {
+        if (iface == this.getClass()) {
             return (T) this;
         }
+        
+        if (!(wrapper instanceof WrapperProxy)) {
+            if (iface.isInstance(wrapper)) {
+                return (T) wrapper;
+            }
+        }
+
 
         return wrapper.unwrap(iface);
     }

src/main/java/com/alibaba/druid/proxy/jdbc/ConnectionProxyImpl.java

@@ -375,10 +375,10 @@ public class ConnectionProxyImpl extends WrapperProxyImpl implements ConnectionP
 
     @SuppressWarnings("unchecked")
     public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == Connection.class) {
-            return (T) connection;
+        if (iface == this.getClass() || iface == ConnectionProxy.class) {
+            return (T) this;
         }
-
+        
         return super.unwrap(iface);
     }
 

src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java

@@ -473,13 +473,4 @@ public class PreparedStatementProxyImpl extends StatementProxyImpl implements Pr
     public String getLastExecuteSql() {
         return this.sql;
     }
-
-    @SuppressWarnings("unchecked")
-    public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == PreparedStatement.class) {
-            return (T) statement;
-        }
-
-        return super.unwrap(iface);
-    }
 }

src/main/java/com/alibaba/druid/proxy/jdbc/ResultSetProxyImpl.java

@@ -1071,15 +1071,6 @@ public class ResultSetProxyImpl extends WrapperProxyImpl implements ResultSetPro
         throw new SQLFeatureNotSupportedException();
     }
 
-    @SuppressWarnings("unchecked")
-    public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == ResultSet.class) {
-            return (T) resultSet;
-        }
-
-        return super.unwrap(iface);
-    }
-
     public int getCloseCount() {
         return closeCount;
     }

src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java

@@ -393,15 +393,6 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro
         throw new SQLFeatureNotSupportedException();
     }
 
-    @SuppressWarnings("unchecked")
-    public <T> T unwrap(Class<T> iface) throws SQLException {
-        if (iface == Statement.class) {
-            return (T) statement;
-        }
-
-        return super.unwrap(iface);
-    }
-
     @Override
     public Map<Integer, JdbcParameter> getParameters() {
         return Collections.emptyMap();

src/main/java/com/alibaba/druid/proxy/jdbc/WrapperProxyImpl.java

@@ -54,7 +54,7 @@ public abstract class WrapperProxyImpl implements WrapperProxy {
             return false;
         }
 
-        if (iface.isInstance(this)) {
+        if (iface == this.getClass()) {
             return true;
         }
 
@@ -67,8 +67,8 @@ public abstract class WrapperProxyImpl implements WrapperProxy {
         if (iface == null) {
             return null;
         }
-
-        if (iface.isInstance(this)) {
+        
+        if (iface == this.getClass()) {
             return (T) this;
         }