diff --git a/src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java b/src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java index 4120e21de2..047cc463c5 100644 --- a/src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java +++ b/src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java @@ -573,4 +573,21 @@ public class PreparedStatementProxyImpl extends StatementProxyImpl implements Pr public String getLastExecuteSql() { return this.sql; } + + @SuppressWarnings("unchecked") + public T unwrap(Class iface) throws SQLException { + if (iface == PreparedStatementProxy.class) { + return (T) this; + } + + return super.unwrap(iface); + } + + public boolean isWrapperFor(Class iface) throws SQLException { + if (iface == PreparedStatementProxy.class) { + return true; + } + + return super.isWrapperFor(iface); + } } diff --git a/src/main/java/com/alibaba/druid/proxy/jdbc/ResultSetProxyImpl.java b/src/main/java/com/alibaba/druid/proxy/jdbc/ResultSetProxyImpl.java index 6ede757e8f..295e56d634 100644 --- a/src/main/java/com/alibaba/druid/proxy/jdbc/ResultSetProxyImpl.java +++ b/src/main/java/com/alibaba/druid/proxy/jdbc/ResultSetProxyImpl.java @@ -1594,4 +1594,20 @@ public class ResultSetProxyImpl extends WrapperProxyImpl implements ResultSetPro return openReaderCount; } + @SuppressWarnings("unchecked") + public T unwrap(Class iface) throws SQLException { + if (iface == ResultSetProxy.class || iface == ResultSetProxyImpl.class) { + return (T) this; + } + + return super.unwrap(iface); + } + + public boolean isWrapperFor(Class iface) throws SQLException { + if (iface == ResultSetProxy.class || iface == ResultSetProxyImpl.class) { + return true; + } + + return super.isWrapperFor(iface); + } } diff --git a/src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java b/src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java index 2b81b702f2..c983019d18 100644 --- a/src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java +++ b/src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java @@ -549,4 +549,21 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro public boolean isFirstResultSet() { return firstResultSet; } + + @SuppressWarnings("unchecked") + public T unwrap(Class iface) throws SQLException { + if (iface == StatementProxy.class) { + return (T) this; + } + + return super.unwrap(iface); + } + + public boolean isWrapperFor(Class iface) throws SQLException { + if (iface == StatementProxy.class) { + return true; + } + + return super.isWrapperFor(iface); + } } diff --git a/src/test/java/com/alibaba/druid/bvt/pool/PreparedStatementProxyImplTest.java b/src/test/java/com/alibaba/druid/bvt/pool/PreparedStatementProxyImplTest.java index edd3483db2..5e0da955d0 100644 --- a/src/test/java/com/alibaba/druid/bvt/pool/PreparedStatementProxyImplTest.java +++ b/src/test/java/com/alibaba/druid/bvt/pool/PreparedStatementProxyImplTest.java @@ -10,7 +10,7 @@ import junit.framework.TestCase; import org.junit.Assert; import com.alibaba.druid.pool.DruidDataSource; -import com.alibaba.druid.proxy.jdbc.PreparedStatementProxyImpl; +import com.alibaba.druid.proxy.jdbc.PreparedStatementProxy; public class PreparedStatementProxyImplTest extends TestCase { @@ -50,7 +50,7 @@ public class PreparedStatementProxyImplTest extends TestCase { stmt.execute(); - PreparedStatementProxyImpl stmtProxy = stmt.unwrap(PreparedStatementProxyImpl.class); + PreparedStatementProxy stmtProxy = stmt.unwrap(PreparedStatementProxy.class); Assert.assertNotNull(stmtProxy); Assert.assertEquals(Types.TINYINT, stmtProxy.getParameters().get(1).getSqlType()); diff --git a/src/test/java/com/alibaba/druid/bvt/pool/UnwrapTest.java b/src/test/java/com/alibaba/druid/bvt/pool/UnwrapTest.java new file mode 100644 index 0000000000..ac4439c8d7 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/pool/UnwrapTest.java @@ -0,0 +1,95 @@ +package com.alibaba.druid.bvt.pool; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; + +import junit.framework.TestCase; + +import org.junit.Assert; + +import com.alibaba.druid.mock.MockConnection; +import com.alibaba.druid.mock.MockPreparedStatement; +import com.alibaba.druid.mock.MockResultSet; +import com.alibaba.druid.pool.DruidDataSource; +import com.alibaba.druid.pool.DruidPooledConnection; +import com.alibaba.druid.proxy.jdbc.PreparedStatementProxy; +import com.alibaba.druid.proxy.jdbc.PreparedStatementProxyImpl; +import com.alibaba.druid.proxy.jdbc.ResultSetProxy; +import com.alibaba.druid.proxy.jdbc.ResultSetProxyImpl; +import com.alibaba.druid.proxy.jdbc.StatementProxy; + +public class UnwrapTest extends TestCase { + + private DruidDataSource dataSource; + + protected void setUp() throws Exception { + + dataSource = new DruidDataSource(); + dataSource.setUrl("jdbc:mock:xxx"); + dataSource.setTestOnBorrow(true); + dataSource.setInitialSize(1); + dataSource.setValidationQuery("select 1"); + dataSource.setValidationQueryTimeout(10); + dataSource.setQueryTimeout(100); + + dataSource.setFilters("log4j"); + } + + protected void tearDown() throws Exception { + dataSource.close(); + } + + public void test_unwrap() throws Exception { + Connection conn = dataSource.getConnection(); + + Assert.assertTrue(conn.isWrapperFor(DruidPooledConnection.class)); + Assert.assertNotNull(conn.unwrap(DruidPooledConnection.class)); + Assert.assertSame(conn, conn.unwrap(DruidPooledConnection.class)); + + Assert.assertTrue(conn.isWrapperFor(MockConnection.class)); + Assert.assertNotNull(conn.unwrap(MockConnection.class)); + + PreparedStatement stmt = conn.prepareStatement("select ?"); + + Assert.assertTrue(stmt.isWrapperFor(Statement.class)); + Assert.assertNotNull(stmt.unwrap(Statement.class)); + + Assert.assertTrue(stmt.isWrapperFor(PreparedStatement.class)); + Assert.assertNotNull(stmt.unwrap(PreparedStatement.class)); + + Assert.assertTrue(stmt.isWrapperFor(StatementProxy.class)); + Assert.assertNotNull(stmt.unwrap(StatementProxy.class)); + + Assert.assertTrue(stmt.isWrapperFor(PreparedStatementProxy.class)); + Assert.assertNotNull(stmt.unwrap(PreparedStatementProxy.class)); + + Assert.assertTrue(stmt.isWrapperFor(PreparedStatementProxyImpl.class)); + Assert.assertNotNull(stmt.unwrap(PreparedStatementProxyImpl.class)); + + Assert.assertTrue(stmt.isWrapperFor(MockPreparedStatement.class)); + Assert.assertNotNull(stmt.unwrap(MockPreparedStatement.class)); + + stmt.setObject(1, "aaa"); + ResultSet rs = stmt.executeQuery(); + + Assert.assertTrue(rs.isWrapperFor(ResultSet.class)); + Assert.assertNotNull(rs.unwrap(ResultSet.class)); + + Assert.assertTrue(rs.isWrapperFor(ResultSetProxy.class)); + Assert.assertNotNull(rs.unwrap(ResultSetProxy.class)); + + Assert.assertTrue(rs.isWrapperFor(ResultSetProxyImpl.class)); + Assert.assertNotNull(rs.unwrap(ResultSetProxyImpl.class)); + + Assert.assertTrue(rs.isWrapperFor(MockResultSet.class)); + Assert.assertNotNull(rs.unwrap(MockResultSet.class)); + + rs.close(); + + stmt.close(); + + conn.close(); + } +} diff --git a/src/test/java/com/alibaba/druid/bvt/pool/UnwrapTest2.java b/src/test/java/com/alibaba/druid/bvt/pool/UnwrapTest2.java new file mode 100644 index 0000000000..138e36efe5 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/pool/UnwrapTest2.java @@ -0,0 +1,86 @@ +package com.alibaba.druid.bvt.pool; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; + +import junit.framework.TestCase; + +import org.junit.Assert; + +import com.alibaba.druid.mock.MockConnection; +import com.alibaba.druid.mock.MockResultSet; +import com.alibaba.druid.mock.MockStatement; +import com.alibaba.druid.pool.DruidDataSource; +import com.alibaba.druid.pool.DruidPooledConnection; +import com.alibaba.druid.proxy.jdbc.ResultSetProxy; +import com.alibaba.druid.proxy.jdbc.ResultSetProxyImpl; +import com.alibaba.druid.proxy.jdbc.StatementProxy; +import com.alibaba.druid.proxy.jdbc.StatementProxyImpl; + +public class UnwrapTest2 extends TestCase { + + private DruidDataSource dataSource; + + protected void setUp() throws Exception { + + dataSource = new DruidDataSource(); + dataSource.setUrl("jdbc:mock:xxx"); + dataSource.setTestOnBorrow(true); + dataSource.setInitialSize(1); + dataSource.setValidationQuery("select 1"); + dataSource.setValidationQueryTimeout(10); + dataSource.setQueryTimeout(100); + + dataSource.setFilters("log4j"); + } + + protected void tearDown() throws Exception { + dataSource.close(); + } + + public void test_unwrap() throws Exception { + Connection conn = dataSource.getConnection(); + + Assert.assertTrue(conn.isWrapperFor(DruidPooledConnection.class)); + Assert.assertNotNull(conn.unwrap(DruidPooledConnection.class)); + Assert.assertSame(conn, conn.unwrap(DruidPooledConnection.class)); + + Assert.assertTrue(conn.isWrapperFor(MockConnection.class)); + Assert.assertNotNull(conn.unwrap(MockConnection.class)); + + Statement stmt = conn.createStatement(); + + Assert.assertTrue(stmt.isWrapperFor(Statement.class)); + Assert.assertNotNull(stmt.unwrap(Statement.class)); + + Assert.assertTrue(stmt.isWrapperFor(StatementProxy.class)); + Assert.assertNotNull(stmt.unwrap(StatementProxy.class)); + + Assert.assertTrue(stmt.isWrapperFor(StatementProxyImpl.class)); + Assert.assertNotNull(stmt.unwrap(StatementProxyImpl.class)); + + Assert.assertTrue(stmt.isWrapperFor(MockStatement.class)); + Assert.assertNotNull(stmt.unwrap(MockStatement.class)); + + ResultSet rs = stmt.executeQuery("select 1"); + + Assert.assertTrue(rs.isWrapperFor(ResultSet.class)); + Assert.assertNotNull(rs.unwrap(ResultSet.class)); + + Assert.assertTrue(rs.isWrapperFor(ResultSetProxy.class)); + Assert.assertNotNull(rs.unwrap(ResultSetProxy.class)); + + Assert.assertTrue(rs.isWrapperFor(ResultSetProxyImpl.class)); + Assert.assertNotNull(rs.unwrap(ResultSetProxyImpl.class)); + + Assert.assertTrue(rs.isWrapperFor(MockResultSet.class)); + Assert.assertNotNull(rs.unwrap(MockResultSet.class)); + + rs.close(); + + stmt.close(); + + conn.close(); + } +}