bug fixed for duplicate getUpdateCount

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

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);
+}

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);
+    }
+
+}

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<MockConnection>     connections           = new CopyOnWriteArrayList<MockConnection>();
+    private final List<MockConnection>     connections              = new CopyOnWriteArrayList<MockConnection>();
 
-    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<MockConnection> 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());

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);
+}

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);
+    }
+
+}

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

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<String>    batchSqlList;
 
-    protected ArrayList<String>   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;
+    }
 }

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;
+        }
+    }
+}

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();
 
         {