diff --git a/src/main/java/com/alibaba/druid/wall/WallCheckResult.java b/src/main/java/com/alibaba/druid/wall/WallCheckResult.java index 16e8fb2276..c7fc201954 100644 --- a/src/main/java/com/alibaba/druid/wall/WallCheckResult.java +++ b/src/main/java/com/alibaba/druid/wall/WallCheckResult.java @@ -36,22 +36,26 @@ public class WallCheckResult { this(null); } - public WallCheckResult(WallSqlStat sqlStat){ + public WallCheckResult(WallSqlStat sqlStat, List stmtList){ if (sqlStat != null) { tableStats = sqlStat.getTableStats(); violations = sqlStat.getViolations(); functionStats = sqlStat.getFunctionStats(); - statementList = Collections.emptyList(); + statementList = stmtList; syntaxError = sqlStat.isSyntaxError(); } else { tableStats = Collections.emptyMap(); violations = Collections.emptyList(); functionStats = Collections.emptyMap(); - statementList = Collections.emptyList(); + statementList = stmtList; syntaxError = false; } } + public WallCheckResult(WallSqlStat sqlStat){ + this(sqlStat, Collections. emptyList()); + } + public WallCheckResult(List violations, Map tableStats, Map functionStats, List statementList, boolean syntaxError){ @@ -70,7 +74,6 @@ public class WallCheckResult { return statementList; } - public Map getTableStats() { return tableStats; } diff --git a/src/main/java/com/alibaba/druid/wall/WallContext.java b/src/main/java/com/alibaba/druid/wall/WallContext.java index 2a064cb7e5..d8a6ea6f09 100644 --- a/src/main/java/com/alibaba/druid/wall/WallContext.java +++ b/src/main/java/com/alibaba/druid/wall/WallContext.java @@ -22,7 +22,7 @@ public class WallContext { private final static ThreadLocal contextLocal = new ThreadLocal(); - private WallSqlStat sqlState; + private WallSqlStat sqlStat; private Map tableStats; private Map functionStats; private final String dbType; @@ -97,12 +97,12 @@ public class WallContext { contextLocal.set(context); } - public WallSqlStat getSqlState() { - return sqlState; + public WallSqlStat getSqlStat() { + return sqlStat; } - public void setSqlState(WallSqlStat sqlState) { - this.sqlState = sqlState; + public void setSqlStat(WallSqlStat sqlStat) { + this.sqlStat = sqlStat; } public Map getTableStats() { diff --git a/src/main/java/com/alibaba/druid/wall/WallFilter.java b/src/main/java/com/alibaba/druid/wall/WallFilter.java index 85d06c6095..c24b1df845 100644 --- a/src/main/java/com/alibaba/druid/wall/WallFilter.java +++ b/src/main/java/com/alibaba/druid/wall/WallFilter.java @@ -20,6 +20,7 @@ import java.sql.SQLException; import java.sql.Wrapper; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Set; import com.alibaba.druid.DruidRuntimeException; @@ -55,7 +56,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { private volatile boolean logViolation = false; private volatile boolean throwException = true; - public final static String ATTR_CONTEXT = "wall.ctx"; + public final static String ATTR_SQL_STAT = "wall.sqlStat"; @Override public void init(DataSourceProxy dataSource) { @@ -164,7 +165,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public void statement_addBatch(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); chain.statement_addBatch(statement, sql); @@ -180,7 +181,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareStatement(connection, sql); + PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -193,7 +196,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareStatement(connection, sql, autoGeneratedKeys); + PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, autoGeneratedKeys); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -207,7 +212,10 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareStatement(connection, sql, resultSetType, resultSetConcurrency); + PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, resultSetType, + resultSetConcurrency); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -221,8 +229,10 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareStatement(connection, sql, resultSetType, resultSetConcurrency, - resultSetHoldability); + PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, resultSetType, + resultSetConcurrency, resultSetHoldability); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -235,7 +245,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareStatement(connection, sql, columnIndexes); + PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, columnIndexes); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -248,7 +260,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareStatement(connection, sql, columnNames); + PreparedStatementProxy stmt = chain.connection_prepareStatement(connection, sql, columnNames); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -261,7 +275,9 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareCall(connection, sql); + CallableStatementProxy stmt = chain.connection_prepareCall(connection, sql); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -275,7 +291,10 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareCall(connection, sql, resultSetType, resultSetConcurrency); + CallableStatementProxy stmt = chain.connection_prepareCall(connection, sql, resultSetType, + resultSetConcurrency); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -289,8 +308,10 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { WallContext.createIfNotExists(dbType); try { sql = check(sql); - return chain.connection_prepareCall(connection, sql, resultSetType, resultSetConcurrency, - resultSetHoldability); + CallableStatementProxy stmt = chain.connection_prepareCall(connection, sql, resultSetType, + resultSetConcurrency, resultSetHoldability); + setSqlStatAttribute(stmt); + return stmt; } finally { WallContext.clearContext(); } @@ -300,13 +321,16 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); - WallContext originalContext = WallContext.current(); try { - WallContext.create(dbType); + createWallContext(statement); sql = check(sql); - return chain.statement_execute(statement, sql); + boolean firstResult = chain.statement_execute(statement, sql); + if (!firstResult) { + int updateCount = statement.getUpdateCount(); + statExecuteUpdate(updateCount); + } + return firstResult; } finally { if (originalContext != null) { WallContext.setContext(originalContext); @@ -317,10 +341,15 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql, int autoGeneratedKeys) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_execute(statement, sql, autoGeneratedKeys); + boolean firstResult = chain.statement_execute(statement, sql, autoGeneratedKeys); + if (!firstResult) { + int updateCount = statement.getUpdateCount(); + statExecuteUpdate(updateCount); + } + return firstResult; } finally { WallContext.clearContext(); } @@ -329,10 +358,15 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql, int columnIndexes[]) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_execute(statement, sql, columnIndexes); + boolean firstResult = chain.statement_execute(statement, sql, columnIndexes); + if (!firstResult) { + int updateCount = statement.getUpdateCount(); + statExecuteUpdate(updateCount); + } + return firstResult; } finally { WallContext.clearContext(); } @@ -341,10 +375,15 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public boolean statement_execute(FilterChain chain, StatementProxy statement, String sql, String columnNames[]) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_execute(statement, sql, columnNames); + boolean firstResult = chain.statement_execute(statement, sql, columnNames); + if (!firstResult) { + int updateCount = statement.getUpdateCount(); + statExecuteUpdate(updateCount); + } + return firstResult; } finally { WallContext.clearContext(); } @@ -353,7 +392,7 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public ResultSetProxy statement_executeQuery(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); return chain.statement_executeQuery(statement, sql); @@ -364,10 +403,12 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_executeUpdate(statement, sql); + int updateCount = chain.statement_executeUpdate(statement, sql); + statExecuteUpdate(updateCount); + return updateCount; } finally { WallContext.clearContext(); } @@ -376,10 +417,12 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql, int autoGeneratedKeys) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_executeUpdate(statement, sql, autoGeneratedKeys); + int updateCount = chain.statement_executeUpdate(statement, sql, autoGeneratedKeys); + statExecuteUpdate(updateCount); + return updateCount; } finally { WallContext.clearContext(); } @@ -388,33 +431,88 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql, int columnIndexes[]) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_executeUpdate(statement, sql, columnIndexes); + int updateCount = chain.statement_executeUpdate(statement, sql, columnIndexes); + statExecuteUpdate(updateCount); + return updateCount; } finally { WallContext.clearContext(); } } - private void setWallContext(StatementProxy statement) { - String dbType = statement.getConnectionProxy().getDirectDataSource().getDbType(); + public String getDbType(StatementProxy statement) { + return statement.getConnectionProxy().getDirectDataSource().getDbType(); + } + + private WallContext createWallContext(StatementProxy statement) { + String dbType = getDbType(statement); WallContext context = WallContext.create(dbType); - statement.getAttributes().put(ATTR_CONTEXT, context); + return context; } @Override public int statement_executeUpdate(FilterChain chain, StatementProxy statement, String sql, String columnNames[]) throws SQLException { - setWallContext(statement); + createWallContext(statement); try { sql = check(sql); - return chain.statement_executeUpdate(statement, sql, columnNames); + int updateCount = chain.statement_executeUpdate(statement, sql, columnNames); + statExecuteUpdate(updateCount); + return updateCount; } finally { WallContext.clearContext(); } } + public void setSqlStatAttribute(PreparedStatementProxy stmt) { + WallContext context = WallContext.current(); + if (context == null) { + return; + } + + WallSqlStat sqlStat = context.getSqlStat(); + if (sqlStat == null) { + return; + } + + stmt.getAttributes().put(ATTR_SQL_STAT, sqlStat); + } + + public void statExecuteUpdate(int updateCount) { + WallContext context = WallContext.current(); + if (context == null) { + return; + } + + WallSqlStat sqlStat = context.getSqlStat(); + if (sqlStat == null) { + return; + } + + Map sqlTableStats = sqlStat.getTableStats(); + if (sqlTableStats == null) { + return; + } + + for (Map.Entry entry : sqlTableStats.entrySet()) { + String tableName = entry.getKey(); + WallTableStat tableStat = provider.getTableStat(tableName); + if (tableStat == null) { + continue; + } + + WallSqlTableStat sqlTableStat = entry.getValue(); + + if (sqlTableStat.getDeleteCount() > 0) { + tableStat.addDeleteDataCount(updateCount); + } else if (sqlTableStat.getUpdateCount() > 0) { + tableStat.addUpdateDataCount(updateCount); + } + } + } + public String check(String sql) throws SQLException { WallCheckResult checkResult = provider.check(sql); List violations = checkResult.getViolations(); @@ -485,6 +583,36 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { return chain.connection_getMetaData(connection); } + @Override + public void resultSet_close(FilterChain chain, ResultSetProxy resultSet) throws SQLException { + chain.resultSet_close(resultSet); + int fetchRowCount = resultSet.getFetchRowCount(); + + WallSqlStat sqlStat = (WallSqlStat) resultSet.getStatementProxy().getAttributes().get(ATTR_SQL_STAT); + if (sqlStat == null) { + return; + } + + Map sqlTableStats = sqlStat.getTableStats(); + if (sqlTableStats == null) { + return; + } + + for (Map.Entry entry : sqlTableStats.entrySet()) { + String tableName = entry.getKey(); + WallTableStat tableStat = provider.getTableStat(tableName); + if (tableStat == null) { + continue; + } + + WallSqlTableStat sqlTableStat = entry.getValue(); + + if (sqlTableStat.getSelectCount() > 0) { + tableStat.addFetchRowCount(fetchRowCount); + } + } + } + public long getViolationCount() { return this.provider.getViolationCount(); } diff --git a/src/main/java/com/alibaba/druid/wall/WallProvider.java b/src/main/java/com/alibaba/druid/wall/WallProvider.java index cc8681001a..61c20165d8 100644 --- a/src/main/java/com/alibaba/druid/wall/WallProvider.java +++ b/src/main/java/com/alibaba/druid/wall/WallProvider.java @@ -454,7 +454,8 @@ public abstract class WallProvider { } if (sqlStat != null) { - new WallCheckResult(sqlStat); + context.setSqlStat(sqlStat); + return new WallCheckResult(sqlStat, statementList); } return new WallCheckResult(violations, tableStats, functionStats, statementList, syntaxError); diff --git a/src/main/java/com/alibaba/druid/wall/WallSqlStat.java b/src/main/java/com/alibaba/druid/wall/WallSqlStat.java index d431f360f0..8dede70458 100644 --- a/src/main/java/com/alibaba/druid/wall/WallSqlStat.java +++ b/src/main/java/com/alibaba/druid/wall/WallSqlStat.java @@ -70,5 +70,4 @@ public class WallSqlStat { public boolean isSyntaxError() { return syntaxError; } - } diff --git a/src/main/java/com/alibaba/druid/wall/WallTableStat.java b/src/main/java/com/alibaba/druid/wall/WallTableStat.java index 36a25cf2e2..3a455b0f19 100644 --- a/src/main/java/com/alibaba/druid/wall/WallTableStat.java +++ b/src/main/java/com/alibaba/druid/wall/WallTableStat.java @@ -33,6 +33,9 @@ public class WallTableStat { private volatile long alterCount; private volatile long dropCount; private volatile long replaceCount; + private volatile long deleteDataCount; + private volatile long updateDataCount; + private volatile long fetchRowCount; final static AtomicLongFieldUpdater selectCountUpdater = AtomicLongFieldUpdater.newUpdater(WallTableStat.class, "selectCount"); @@ -55,6 +58,14 @@ public class WallTableStat { final static AtomicLongFieldUpdater replaceCountUpdater = AtomicLongFieldUpdater.newUpdater(WallTableStat.class, "replaceCount"); + final static AtomicLongFieldUpdater deleteDataCountUpdater = AtomicLongFieldUpdater.newUpdater(WallTableStat.class, + "deleteDataCount"); + final static AtomicLongFieldUpdater updateDataCountUpdater = AtomicLongFieldUpdater.newUpdater(WallTableStat.class, + "updateDataCount"); + + final static AtomicLongFieldUpdater fetchRowCountUpdater = AtomicLongFieldUpdater.newUpdater(WallTableStat.class, + "fetchRowCount"); + public WallTableStat(){ } @@ -94,11 +105,35 @@ public class WallTableStat { public long getDropCount() { return dropCount; } - + public long getReplaceCount() { return replaceCount; } + public long getDeleteDataCount() { + return this.deleteDataCount; + } + + public void addDeleteDataCount(long delta) { + deleteDataCountUpdater.addAndGet(this, delta); + } + + public long getUpdateDataCount() { + return this.updateDataCount; + } + + public void addUpdateDataCount(long delta) { + updateDataCountUpdater.addAndGet(this, delta); + } + + public long getFetchRowCount() { + return fetchRowCount; + } + + public void addFetchRowCount(long delta) { + fetchRowCountUpdater.addAndGet(this, delta); + } + public void addSqlTableStat(WallSqlTableStat stat) { { long val = stat.getSelectCount(); @@ -202,6 +237,15 @@ public class WallTableStat { if (replaceCount > 0) { map.put("replaceCount", replaceCount); } + if (deleteDataCount > 0) { + map.put("deleteDataCount", deleteDataCount); + } + if (fetchRowCount > 0) { + map.put("fetchRowCount", fetchRowCount); + } + if (updateDataCount > 0) { + map.put("updateDataCount", updateDataCount); + } return map; } } diff --git a/src/main/resources/support/http/resources/wall.html b/src/main/resources/support/http/resources/wall.html index eac8bd835f..1392b54144 100644 --- a/src/main/resources/support/http/resources/wall.html +++ b/src/main/resources/support/http/resources/wall.html @@ -73,6 +73,9 @@ AlterCount DropCount ReplaceCount + DeleteDataCount + UpdateDataCount + FetchRowCount @@ -128,6 +131,9 @@ html += "" + (table.alterCount ? table.alterCount : 0) + ""; html += "" + (table.dropCount ? table.dropCount : 0) + ""; html += "" + (table.replaceCount ? table.replaceCount : 0) + ""; + html += "" + (table.deleteDataCount ? table.deleteDataCount : 0) + ""; + html += "" + (table.updateDataCount ? table.updateDataCount : 0) + ""; + html += "" + (table.fetchRowCount ? table.fetchRowCount : 0) + ""; html += ""; } $("#dataTable1 tbody").html(html); 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 new file mode 100644 index 0000000000..1a74df98b9 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallFilterTest.java @@ -0,0 +1,109 @@ +package com.alibaba.druid.bvt.filter.wall; + +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.pool.DruidDataSource; +import com.alibaba.druid.wall.WallFilter; + +public class WallFilterTest extends TestCase { + + private DruidDataSource dataSource; + private WallFilter wallFilter; + + protected void setUp() throws Exception { + dataSource = new DruidDataSource(); + + dataSource.setUrl("jdbc:h2:mem:wall_test;"); + dataSource.setFilters("wall"); + dataSource.init(); + + wallFilter = (WallFilter) dataSource.getProxyFilters().get(0); + } + + protected void tearDown() throws Exception { + dataSource.close(); + } + + public void test_wallFilter() throws Exception { + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + stmt.execute("CREATE TABLE t (FID INTEGER, FNAME VARCHAR(50))"); + stmt.close(); + conn.close(); + } + Assert.assertEquals(1, wallFilter.getProvider().getTableStat("t").getCreateCount()); + + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + for (int i = 0; i < 10; ++i) { + stmt.execute("INSERT INTO t (FID, FNAME) VALUES (" + i + ", 'a" + i + "')"); + } + stmt.close(); + conn.close(); + } + { + String sql = "SELECT * FROM T"; + + Connection conn = dataSource.getConnection(); + PreparedStatement stmt = conn.prepareStatement(sql); + ResultSet rs = stmt.executeQuery(); + while (rs.next()) { + + } + rs.close(); + stmt.close(); + conn.close(); + } + Assert.assertEquals(10, wallFilter.getProvider().getTableStat("t").getFetchRowCount()); + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + stmt.execute("DELETE from t where FID = 0"); + stmt.close(); + conn.close(); + } + Assert.assertEquals(1, wallFilter.getProvider().getTableStat("t").getDeleteDataCount()); + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + stmt.execute("DELETE from t where FID = 1 OR FID = 2"); + stmt.close(); + conn.close(); + } + Assert.assertEquals(3, wallFilter.getProvider().getTableStat("t").getDeleteDataCount()); + + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + stmt.execute("update t SET fname = 'xx' where FID = 3 OR FID = 4"); + stmt.close(); + conn.close(); + } + Assert.assertEquals(2, wallFilter.getProvider().getTableStat("t").getUpdateDataCount()); + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + stmt.execute("truncate table t"); + stmt.close(); + conn.close(); + } + Assert.assertEquals(1, wallFilter.getProvider().getTableStat("t").getTruncateCount()); + { + Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement(); + stmt.execute("drop table t"); + stmt.close(); + conn.close(); + } + Assert.assertEquals(1, wallFilter.getProvider().getTableStat("t").getDropCount()); + } +}