diff --git a/src/main/java/com/alibaba/druid/wall/WallFilter.java b/src/main/java/com/alibaba/druid/wall/WallFilter.java index c24b1df845..908ef28659 100644 --- a/src/main/java/com/alibaba/druid/wall/WallFilter.java +++ b/src/main/java/com/alibaba/druid/wall/WallFilter.java @@ -178,7 +178,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { public PreparedStatementProxy connection_prepareStatement(FilterChain chain, ConnectionProxy connection, String sql) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql); @@ -193,7 +193,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { public PreparedStatementProxy connection_prepareStatement(FilterChain chain, ConnectionProxy connection, String sql, int autoGeneratedKeys) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, autoGeneratedKeys); @@ -209,7 +209,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { String sql, int resultSetType, int resultSetConcurrency) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, resultSetType, @@ -226,7 +226,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, resultSetType, @@ -242,7 +242,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { public PreparedStatementProxy connection_prepareStatement(FilterChain chain, ConnectionProxy connection, String sql, int[] columnIndexes) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, columnIndexes); @@ -257,7 +257,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { public PreparedStatementProxy connection_prepareStatement(FilterChain chain, ConnectionProxy connection, String sql, String[] columnNames) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, columnNames); @@ -272,7 +272,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { public CallableStatementProxy connection_prepareCall(FilterChain chain, ConnectionProxy connection, String sql) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); CallableStatementProxy stmt = chain.connection_prepareCall(connection, sql); @@ -288,7 +288,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { int resultSetType, int resultSetConcurrency) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); CallableStatementProxy stmt = chain.connection_prepareCall(connection, sql, resultSetType, @@ -305,7 +305,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { String dbType = connection.getDirectDataSource().getDbType(); - WallContext.createIfNotExists(dbType); + WallContext.create(dbType); try { sql = check(sql); CallableStatementProxy stmt = chain.connection_prepareCall(connection, sql, resultSetType, @@ -465,6 +465,33 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.clearContext(); } } + + @Override + public boolean preparedStatement_execute(FilterChain chain, PreparedStatementProxy statement) throws SQLException { + boolean firstResult = chain.preparedStatement_execute(statement); + + if (!firstResult) { + WallSqlStat sqlStat = (WallSqlStat) statement.getAttributes().get(ATTR_SQL_STAT); + int updateCount = statement.getUpdateCount(); + statExecuteUpdate(sqlStat, updateCount); + } + + return firstResult; + } + + @Override + public ResultSetProxy preparedStatement_executeQuery(FilterChain chain, PreparedStatementProxy statement) + throws SQLException { + return chain.preparedStatement_executeQuery(statement); + } + + @Override + public int preparedStatement_executeUpdate(FilterChain chain, PreparedStatementProxy statement) throws SQLException { + int updateCount = chain.preparedStatement_executeUpdate(statement); + statExecuteUpdate(updateCount); + return updateCount; + } + public void setSqlStatAttribute(PreparedStatementProxy stmt) { WallContext context = WallContext.current(); @@ -491,6 +518,10 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { return; } + statExecuteUpdate(sqlStat, updateCount); + } + + private void statExecuteUpdate(WallSqlStat sqlStat, int updateCount) { Map sqlTableStats = sqlStat.getTableStats(); if (sqlTableStats == null) { return; diff --git a/src/main/java/com/alibaba/druid/wall/WallProvider.java b/src/main/java/com/alibaba/druid/wall/WallProvider.java index 61c20165d8..1c243367db 100644 --- a/src/main/java/com/alibaba/druid/wall/WallProvider.java +++ b/src/main/java/com/alibaba/druid/wall/WallProvider.java @@ -473,6 +473,10 @@ public abstract class WallProvider { } recordStats(sqlStat.getTableStats(), sqlStat.getFunctionStats()); + WallContext context = WallContext.current(); + if (context != null) { + context.setSqlStat(sqlStat); + } return new WallCheckResult(sqlStat); } } diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallFilterTest.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallFilterTest.java index 1a74df98b9..fbbd6bee5d 100644 --- a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallFilterTest.java +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallFilterTest.java @@ -80,7 +80,7 @@ public class WallFilterTest extends TestCase { conn.close(); } Assert.assertEquals(3, wallFilter.getProvider().getTableStat("t").getDeleteDataCount()); - + { Connection conn = dataSource.getConnection(); Statement stmt = conn.createStatement(); @@ -89,6 +89,29 @@ public class WallFilterTest extends TestCase { conn.close(); } Assert.assertEquals(2, wallFilter.getProvider().getTableStat("t").getUpdateDataCount()); + + { + Connection conn = dataSource.getConnection(); + PreparedStatement stmt = conn.prepareStatement("update t SET fname = 'xx' where FID = ? OR FID = ?"); + stmt.setInt(1, 3); + stmt.setInt(2, 4); + stmt.execute(); + stmt.close(); + conn.close(); + } + Assert.assertEquals(4, wallFilter.getProvider().getTableStat("t").getUpdateDataCount()); + + { + Connection conn = dataSource.getConnection(); + PreparedStatement stmt = conn.prepareStatement("update t SET fname = 'xx' where FID = ? OR FID = ?"); + stmt.setInt(1, 3); + stmt.setInt(2, 4); + stmt.execute(); + stmt.close(); + conn.close(); + } + Assert.assertEquals(6, wallFilter.getProvider().getTableStat("t").getUpdateDataCount()); + { Connection conn = dataSource.getConnection(); Statement stmt = conn.createStatement();