diff --git a/src/main/java/com/alibaba/druid/mock/MockConnection.java b/src/main/java/com/alibaba/druid/mock/MockConnection.java index 33345a8506..a460852b4a 100644 --- a/src/main/java/com/alibaba/druid/mock/MockConnection.java +++ b/src/main/java/com/alibaba/druid/mock/MockConnection.java @@ -117,7 +117,7 @@ public class MockConnection extends ConnectionBase implements Connection { throw new MockConnectionClosedException(); } - return new MockPreparedStatement(this, sql); + return createMockPreparedStatement(sql); } @Override @@ -204,7 +204,7 @@ public class MockConnection extends ConnectionBase implements Connection { throw new MockConnectionClosedException(); } - MockPreparedStatement stmt = new MockPreparedStatement(this, sql); + MockPreparedStatement stmt = createMockPreparedStatement(sql); stmt.setResultSetType(resultSetType); stmt.setResultSetConcurrency(resultSetConcurrency); @@ -289,7 +289,7 @@ public class MockConnection extends ConnectionBase implements Connection { throw new MockConnectionClosedException(); } - MockPreparedStatement stmt = new MockPreparedStatement(this, sql); + MockPreparedStatement stmt = createMockPreparedStatement(sql); stmt.setResultSetType(resultSetType); stmt.setResultSetConcurrency(resultSetConcurrency); @@ -298,6 +298,16 @@ public class MockConnection extends ConnectionBase implements Connection { return stmt; } + private MockPreparedStatement createMockPreparedStatement(String sql) { + if (driver != null) { + MockPreparedStatementFactory preparedStatementFactory = driver.getPreparedStatementFactory(); + if (preparedStatementFactory != null) { + return preparedStatementFactory.createMockPreparedStatement(this, sql); + } + } + return new MockPreparedStatement(this, sql); + } + @Override public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { @@ -320,7 +330,7 @@ public class MockConnection extends ConnectionBase implements Connection { throw new MockConnectionClosedException(); } - return new MockPreparedStatement(this, sql); + return createMockPreparedStatement(sql); } @Override @@ -329,7 +339,7 @@ public class MockConnection extends ConnectionBase implements Connection { throw new MockConnectionClosedException(); } - return new MockPreparedStatement(this, sql); + return createMockPreparedStatement(sql); } @Override @@ -338,7 +348,7 @@ public class MockConnection extends ConnectionBase implements Connection { throw new MockConnectionClosedException(); } - return new MockPreparedStatement(this, sql); + return createMockPreparedStatement(sql); } @Override diff --git a/src/main/java/com/alibaba/druid/mock/MockConnectionFactory.java b/src/main/java/com/alibaba/druid/mock/MockConnectionFactory.java new file mode 100644 index 0000000000..d9c1d35b4c --- /dev/null +++ b/src/main/java/com/alibaba/druid/mock/MockConnectionFactory.java @@ -0,0 +1,8 @@ +package com.alibaba.druid.mock; + +import java.util.Properties; + +public interface MockConnectionFactory { + + MockConnection createConnection(MockDriver driver, String url, Properties connectProperties); +} diff --git a/src/main/java/com/alibaba/druid/mock/MockConnectionFactoryImpl.java b/src/main/java/com/alibaba/druid/mock/MockConnectionFactoryImpl.java new file mode 100644 index 0000000000..0726315284 --- /dev/null +++ b/src/main/java/com/alibaba/druid/mock/MockConnectionFactoryImpl.java @@ -0,0 +1,12 @@ +package com.alibaba.druid.mock; + +import java.util.Properties; + +public class MockConnectionFactoryImpl implements MockConnectionFactory { + + @Override + public MockConnection createConnection(MockDriver driver, String url, Properties connectProperties) { + return new MockConnection(driver, url, connectProperties); + } + +} diff --git a/src/main/java/com/alibaba/druid/mock/MockDriver.java b/src/main/java/com/alibaba/druid/mock/MockDriver.java index 1398071f6d..44278c075b 100644 --- a/src/main/java/com/alibaba/druid/mock/MockDriver.java +++ b/src/main/java/com/alibaba/druid/mock/MockDriver.java @@ -43,34 +43,53 @@ import com.alibaba.druid.support.logging.LogFactory; public class MockDriver implements Driver, MockDriverMBean { - private final static Log LOG = LogFactory.getLog(MockDriver.class); + private final static Log LOG = LogFactory.getLog(MockDriver.class); - public final static MockExecuteHandler DEFAULT_HANDLER = new MySqlMockExecuteHandlerImpl(); + public final static MockExecuteHandler DEFAULT_HANDLER = new MySqlMockExecuteHandlerImpl(); - private String prefix = "jdbc:fake:"; - private String mockPrefix = "jdbc:mock:"; + private String prefix = "jdbc:fake:"; + private String mockPrefix = "jdbc:mock:"; - private MockExecuteHandler executeHandler = DEFAULT_HANDLER; + private MockExecuteHandler executeHandler = DEFAULT_HANDLER; - public final static MockDriver instance = new MockDriver(); + public final static MockDriver instance = new MockDriver(); - private final AtomicLong connectCount = new AtomicLong(); - private final AtomicLong connectionCloseCount = new AtomicLong(); + private final AtomicLong connectCount = new AtomicLong(); + private final AtomicLong connectionCloseCount = new AtomicLong(); - private final AtomicLong connectionIdSeed = new AtomicLong(1000L); + private final AtomicLong connectionIdSeed = new AtomicLong(1000L); - private final List connections = new CopyOnWriteArrayList(); + private final List connections = new CopyOnWriteArrayList(); - private long idleTimeCount = 1000 * 60 * 3; + private long idleTimeCount = 1000 * 60 * 3; - private boolean logExecuteQueryEnable = true; + private boolean logExecuteQueryEnable = true; - private final static String MBEAN_NAME = "com.alibaba.druid:type=MockDriver"; + private final static String MBEAN_NAME = "com.alibaba.druid:type=MockDriver"; + + private MockConnectionFactory connectionFactory = new MockConnectionFactoryImpl(); + private MockPreparedStatementFactory preparedStatementFactory = new MockPreparedStatementFactoryImpl(); static { registerDriver(instance); } + public MockConnectionFactory getConnectionFactory() { + return connectionFactory; + } + + public void setConnectionFactory(MockConnectionFactory connectionFactory) { + this.connectionFactory = connectionFactory; + } + + public MockPreparedStatementFactory getPreparedStatementFactory() { + return preparedStatementFactory; + } + + public void setPreparedStatementFactory(MockPreparedStatementFactory preparedStatementFactory) { + this.preparedStatementFactory = preparedStatementFactory; + } + public boolean isLogExecuteQueryEnable() { return logExecuteQueryEnable; } @@ -97,11 +116,11 @@ public class MockDriver implements Driver, MockDriverMBean { conn.close(); } } - + public int getConnectionsSize() { return this.connections.size(); } - + public List getConnections() { return connections; } @@ -154,7 +173,7 @@ public class MockDriver implements Driver, MockDriverMBean { public void setExecuteHandler(MockExecuteHandler executeHandler) { this.executeHandler = executeHandler; } - + @Override public Connection connect(String url, Properties info) throws SQLException { if (!acceptsURL(url)) { @@ -173,7 +192,7 @@ public class MockDriver implements Driver, MockDriverMBean { } } - MockConnection conn = new MockConnection(this, url, info); + MockConnection conn = connectionFactory.createConnection(this, url, info); if (LOG.isDebugEnabled()) { LOG.debug("connect, url " + url + ", id " + conn.getId()); diff --git a/src/main/java/com/alibaba/druid/mock/MockPreparedStatementFactory.java b/src/main/java/com/alibaba/druid/mock/MockPreparedStatementFactory.java new file mode 100644 index 0000000000..de980a19a9 --- /dev/null +++ b/src/main/java/com/alibaba/druid/mock/MockPreparedStatementFactory.java @@ -0,0 +1,6 @@ +package com.alibaba.druid.mock; + + +public interface MockPreparedStatementFactory { + MockPreparedStatement createMockPreparedStatement(MockConnection conn, String sql); +} diff --git a/src/main/java/com/alibaba/druid/mock/MockPreparedStatementFactoryImpl.java b/src/main/java/com/alibaba/druid/mock/MockPreparedStatementFactoryImpl.java new file mode 100644 index 0000000000..d3329a122a --- /dev/null +++ b/src/main/java/com/alibaba/druid/mock/MockPreparedStatementFactoryImpl.java @@ -0,0 +1,11 @@ +package com.alibaba.druid.mock; + + +public class MockPreparedStatementFactoryImpl implements MockPreparedStatementFactory { + + @Override + public MockPreparedStatement createMockPreparedStatement(MockConnection conn, String sql) { + return new MockPreparedStatement(conn, sql); + } + +} 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 21d116a51c..7e9984094c 100644 --- a/src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java +++ b/src/main/java/com/alibaba/druid/proxy/jdbc/PreparedStatementProxyImpl.java @@ -85,6 +85,7 @@ public class PreparedStatementProxyImpl extends StatementProxyImpl implements Pr @Override public boolean execute() throws SQLException { + updateCount = null; lastExecuteType = StatementExecuteType.Execute; firstResultSet = createChain().preparedStatement_execute(this); return firstResultSet; @@ -92,6 +93,7 @@ public class PreparedStatementProxyImpl extends StatementProxyImpl implements Pr @Override public ResultSet executeQuery() throws SQLException { + updateCount = null; lastExecuteType = StatementExecuteType.ExecuteQuery; return createChain().preparedStatement_executeQuery(this); } @@ -99,7 +101,8 @@ public class PreparedStatementProxyImpl extends StatementProxyImpl implements Pr @Override public int executeUpdate() throws SQLException { lastExecuteType = StatementExecuteType.ExecuteUpdate; - return createChain().preparedStatement_executeUpdate(this); + updateCount = createChain().preparedStatement_executeUpdate(this); + return updateCount; } @Override 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 64ed278f70..1908704e2d 100644 --- a/src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java +++ b/src/main/java/com/alibaba/druid/proxy/jdbc/StatementProxyImpl.java @@ -35,18 +35,20 @@ import com.alibaba.druid.stat.JdbcSqlStat; */ public class StatementProxyImpl extends WrapperProxyImpl implements StatementProxy { - private final ConnectionProxy connection; - private final Statement statement; + private final ConnectionProxy connection; + private final Statement statement; - protected String lastExecuteSql; - private long lastExecuteTimeNano; - protected JdbcSqlStat sqlStat; - protected boolean firstResultSet; + protected String lastExecuteSql; + private long lastExecuteTimeNano; + protected JdbcSqlStat sqlStat; + protected boolean firstResultSet; + + protected ArrayList batchSqlList; - protected ArrayList batchSqlList; - protected StatementExecuteType lastExecuteType; + protected Integer updateCount = null; + public StatementProxyImpl(ConnectionProxy connection, Statement statement, long id){ super(statement, id); this.connection = connection; @@ -102,15 +104,17 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro @Override public boolean execute(String sql) throws SQLException { + updateCount = null; lastExecuteSql = sql; lastExecuteType = StatementExecuteType.Execute; - + firstResultSet = createChain().statement_execute(this, sql); return firstResultSet; } @Override public boolean execute(String sql, int autoGeneratedKeys) throws SQLException { + updateCount = null; lastExecuteSql = sql; lastExecuteType = StatementExecuteType.Execute; firstResultSet = createChain().statement_execute(this, sql, autoGeneratedKeys); @@ -119,6 +123,7 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro @Override public boolean execute(String sql, int[] columnIndexes) throws SQLException { + updateCount = null; lastExecuteSql = sql; lastExecuteType = StatementExecuteType.Execute; firstResultSet = createChain().statement_execute(this, sql, columnIndexes); @@ -127,6 +132,7 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro @Override public boolean execute(String sql, String[] columnNames) throws SQLException { + updateCount = null; lastExecuteSql = sql; lastExecuteType = StatementExecuteType.Execute; firstResultSet = createChain().statement_execute(this, sql, columnNames); @@ -135,12 +141,19 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro @Override public int[] executeBatch() throws SQLException { - lastExecuteType = StatementExecuteType.ExecuteBatch; - return createChain().statement_executeBatch(this); + lastExecuteType = StatementExecuteType.ExecuteBatch; + int[] updateCounts = createChain().statement_executeBatch(this); + + if (updateCounts != null && updateCounts.length == 1) { + updateCount = updateCounts[0]; + } + + return updateCounts; } @Override public ResultSet executeQuery(String sql) throws SQLException { + updateCount = null; lastExecuteSql = sql; lastExecuteType = StatementExecuteType.ExecuteQuery; return createChain().statement_executeQuery(this, sql); @@ -150,28 +163,32 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro public int executeUpdate(String sql) throws SQLException { lastExecuteSql = sql; lastExecuteType = StatementExecuteType.ExecuteUpdate; - return createChain().statement_executeUpdate(this, sql); + updateCount = createChain().statement_executeUpdate(this, sql); + return updateCount; } @Override public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException { lastExecuteSql = sql; lastExecuteType = StatementExecuteType.ExecuteUpdate; - return createChain().statement_executeUpdate(this, sql, autoGeneratedKeys); + updateCount = createChain().statement_executeUpdate(this, sql, autoGeneratedKeys); + return updateCount; } @Override public int executeUpdate(String sql, int[] columnIndexes) throws SQLException { lastExecuteSql = sql; lastExecuteType = StatementExecuteType.ExecuteUpdate; - return createChain().statement_executeUpdate(this, sql, columnIndexes); + updateCount = createChain().statement_executeUpdate(this, sql, columnIndexes); + return updateCount; } @Override public int executeUpdate(String sql, String[] columnNames) throws SQLException { lastExecuteSql = sql; lastExecuteType = StatementExecuteType.ExecuteUpdate; - return createChain().statement_executeUpdate(this, sql, columnNames); + updateCount = createChain().statement_executeUpdate(this, sql, columnNames); + return updateCount; } @Override @@ -206,11 +223,13 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro @Override public boolean getMoreResults() throws SQLException { + updateCount = null; return createChain().statement_getMoreResults(this); } @Override public boolean getMoreResults(int current) throws SQLException { + updateCount = null; return createChain().statement_getMoreResults(this, current); } @@ -239,9 +258,13 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro return createChain().statement_getResultSetType(this); } + //bug fixed for oracle @Override public int getUpdateCount() throws SQLException { - return createChain().statement_getUpdateCount(this); + if (updateCount == null) { + updateCount = createChain().statement_getUpdateCount(this); + } + return updateCount; } @Override @@ -356,18 +379,18 @@ public class StatementProxyImpl extends WrapperProxyImpl implements StatementPro } public long getLastExecuteTimeNano() { - return lastExecuteTimeNano; - } + return lastExecuteTimeNano; + } - public void setLastExecuteTimeNano(long lastExecuteTimeNano) { - this.lastExecuteTimeNano = lastExecuteTimeNano; - } + public void setLastExecuteTimeNano(long lastExecuteTimeNano) { + this.lastExecuteTimeNano = lastExecuteTimeNano; + } - public StatementExecuteType getLastExecuteType() { - return lastExecuteType; - } + public StatementExecuteType getLastExecuteType() { + return lastExecuteType; + } - public boolean isFirstResultSet() { - return firstResultSet; - } + public boolean isFirstResultSet() { + return firstResultSet; + } } diff --git a/src/test/java/com/alibaba/druid/bvt/pool/TestGetUpdateCount.java b/src/test/java/com/alibaba/druid/bvt/pool/TestGetUpdateCount.java new file mode 100644 index 0000000000..00f5a28ad0 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/pool/TestGetUpdateCount.java @@ -0,0 +1,104 @@ +package com.alibaba.druid.bvt.pool; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; + +import junit.framework.Assert; +import junit.framework.TestCase; + +import com.alibaba.druid.mock.MockConnection; +import com.alibaba.druid.mock.MockDriver; +import com.alibaba.druid.mock.MockPreparedStatement; +import com.alibaba.druid.mock.MockPreparedStatementFactory; +import com.alibaba.druid.pool.DruidDataSource; +import com.alibaba.druid.util.JdbcUtils; + +public class TestGetUpdateCount extends TestCase { + + private DruidDataSource dataSource; + private MockDriver driver; + + protected void setUp() throws Exception { + driver = new MockDriver(); + driver.setPreparedStatementFactory(new MockPreparedStatementFactory() { + + @Override + public MockPreparedStatement createMockPreparedStatement(MockConnection conn, String sql) { + return new MyPreparedStatement(conn, sql); + } + + }); + + // ///////////////////////////// + + dataSource = new DruidDataSource(); + + dataSource.setDriver(driver); + dataSource.setUrl("jdbc:mock:xx"); + dataSource.setFilters("stat"); + } + + protected void tearDown() throws Exception { + JdbcUtils.close(dataSource); + } + + public void test_executeQuery() throws Exception { + Connection conn = dataSource.getConnection(); + + PreparedStatement stmt = conn.prepareStatement("select ?"); + + MyPreparedStatement myStmt = stmt.unwrap(MyPreparedStatement.class); + + Assert.assertNull(myStmt.updateCount); + + stmt.setString(1, "xxx"); + ResultSet rs = stmt.executeQuery(); + + Assert.assertNull(myStmt.updateCount); + + rs.close(); + stmt.close(); + conn.close(); + } + + public void test_execute() throws Exception { + Connection conn = dataSource.getConnection(); + + PreparedStatement stmt = conn.prepareStatement("update t set id = ?"); + + MyPreparedStatement myStmt = stmt.unwrap(MyPreparedStatement.class); + + Assert.assertNull(myStmt.updateCount); + + stmt.setString(1, "xxx"); + stmt.execute(); + + Assert.assertNotNull(myStmt.updateCount); + + Assert.assertEquals(1, stmt.getUpdateCount()); + + stmt.close(); + conn.close(); + } + + public static class MyPreparedStatement extends MockPreparedStatement { + + Integer updateCount = null; + + public MyPreparedStatement(MockConnection conn, String sql){ + super(conn, sql); + } + + @Override + public int getUpdateCount() throws SQLException { + if (updateCount != null) { + throw new SQLException("illegal state"); + } + + updateCount = 1; + return updateCount; + } + } +} diff --git a/src/test/java/com/alibaba/druid/bvt/pool/basic/PoolableStatementTest2.java b/src/test/java/com/alibaba/druid/bvt/pool/basic/PoolableStatementTest2.java index 0870aeba34..0504815570 100644 --- a/src/test/java/com/alibaba/druid/bvt/pool/basic/PoolableStatementTest2.java +++ b/src/test/java/com/alibaba/druid/bvt/pool/basic/PoolableStatementTest2.java @@ -432,6 +432,7 @@ public class PoolableStatementTest2 extends TestCase { Statement stmt = conn.createStatement(); stmt.getUpdateCount(); + stmt.executeQuery("select 1"); ((DruidPooledStatement) stmt).getStatement().close(); {