From 97f5c7e9ef11832a97741129502ae46c20c9d274 Mon Sep 17 00:00:00 2001
From: wenshao <szujobs@gmail.com>
Date: Sun, 7 Apr 2013 06:45:31 +0800
Subject: [PATCH] improve wall filter

---
 .../com/alibaba/druid/wall/WallFilter.java    | 34 +++++++++++++++----
 .../druid/bvt/filter/wall/WallFilterTest.java | 15 ++++++++
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/alibaba/druid/wall/WallFilter.java b/src/main/java/com/alibaba/druid/wall/WallFilter.java
index 908ef28659..9fd1c2ecb1 100644
--- a/src/main/java/com/alibaba/druid/wall/WallFilter.java
+++ b/src/main/java/com/alibaba/druid/wall/WallFilter.java
@@ -389,6 +389,21 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
         }
     }
 
+    @Override
+    public int[] statement_executeBatch(FilterChain chain, StatementProxy statement) throws SQLException {
+        WallSqlStat sqlStat = (WallSqlStat) statement.getAttributes().get(ATTR_SQL_STAT);
+        try {
+            int[] updateCounts = chain.statement_executeBatch(statement);
+            for (int i = 0; i < updateCounts.length; ++i) {
+                int updateCount = updateCounts[i];
+                statExecuteUpdate(sqlStat, updateCount);
+            }
+            return updateCounts;
+        } finally {
+            WallContext.clearContext();
+        }
+    }
+
     @Override
     public ResultSetProxy statement_executeQuery(FilterChain chain, StatementProxy statement, String sql)
                                                                                                          throws SQLException {
@@ -465,17 +480,17 @@ 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;
     }
 
@@ -492,7 +507,6 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
         return updateCount;
     }
 
-
     public void setSqlStatAttribute(PreparedStatementProxy stmt) {
         WallContext context = WallContext.current();
         if (context == null) {
@@ -518,10 +532,16 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
             return;
         }
 
-        statExecuteUpdate(sqlStat, updateCount);
+        if (updateCount > 0) {
+            statExecuteUpdate(sqlStat, updateCount);
+        }
     }
 
     private void statExecuteUpdate(WallSqlStat sqlStat, int updateCount) {
+        if (sqlStat == null) {
+            return;
+        }
+        
         Map<String, WallSqlTableStat> sqlTableStats = sqlStat.getTableStats();
         if (sqlTableStats == null) {
             return;
@@ -538,7 +558,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
 
             if (sqlTableStat.getDeleteCount() > 0) {
                 tableStat.addDeleteDataCount(updateCount);
-            } else  if (sqlTableStat.getUpdateCount() > 0) {
+            } else if (sqlTableStat.getUpdateCount() > 0) {
                 tableStat.addUpdateDataCount(updateCount);
             }
         }
@@ -628,7 +648,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean {
         if (sqlTableStats == null) {
             return;
         }
-        
+
         for (Map.Entry<String, WallSqlTableStat> entry : sqlTableStats.entrySet()) {
             String tableName = entry.getKey();
             WallTableStat tableStat = provider.getTableStat(tableName);
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 fbbd6bee5d..48849efe25 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
@@ -111,6 +111,21 @@ public class WallFilterTest extends TestCase {
             conn.close();
         }
         Assert.assertEquals(6, wallFilter.getProvider().getTableStat("t").getUpdateDataCount());
+        {
+            Connection conn = dataSource.getConnection();
+            PreparedStatement stmt = conn.prepareStatement("update t SET fname = 'xx' where FID = ?");
+            
+            stmt.setInt(1, 3);
+            stmt.addBatch();
+            
+            stmt.setInt(1, 4);
+            stmt.addBatch();
+            
+            stmt.executeBatch();
+            stmt.close();
+            conn.close();
+        }
+        Assert.assertEquals(8, wallFilter.getProvider().getTableStat("t").getUpdateDataCount());
 
         {
             Connection conn = dataSource.getConnection();