From 23179b0e5cd33a6758d52425751f0350e94ac5d7 Mon Sep 17 00:00:00 2001 From: wenshao Date: Sun, 14 May 2017 02:13:24 +0800 Subject: [PATCH] improved PageUtils.limit & WallFilter support selectLimit. for issue #1369 --- .../com/alibaba/druid/sql/PagerUtils.java | 161 +++++++++++++----- .../druid/sql/ast/statement/SQLSelect.java | 11 ++ .../com/alibaba/druid/wall/WallConfig.java | 10 ++ .../druid/wall/spi/DB2WallVisitor.java | 24 +-- .../druid/wall/spi/MySqlWallVisitor.java | 25 +-- .../druid/wall/spi/OracleWallVisitor.java | 24 +-- .../alibaba/druid/wall/spi/PGWallVisitor.java | 45 +++-- .../druid/wall/spi/SQLServerWallVisitor.java | 24 +-- .../bvt/filter/wall/WallSelectLimitTest.java | 83 +++++++++ .../filter/wall/WallSelectLimitTest_2.java | 87 ++++++++++ 10 files changed, 376 insertions(+), 118 deletions(-) create mode 100644 src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest.java create mode 100644 src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest_2.java diff --git a/src/main/java/com/alibaba/druid/sql/PagerUtils.java b/src/main/java/com/alibaba/druid/sql/PagerUtils.java index 0fced4029d..208b2f56d1 100644 --- a/src/main/java/com/alibaba/druid/sql/PagerUtils.java +++ b/src/main/java/com/alibaba/druid/sql/PagerUtils.java @@ -74,57 +74,74 @@ public class PagerUtils { } public static String limit(SQLSelect select, String dbType, int offset, int count) { + limit(select, dbType, offset, count, false); + + return SQLUtils.toSQLString(select, dbType); + } + + public static boolean limit(SQLSelect select, String dbType, int offset, int count, boolean check) { SQLSelectQuery query = select.getQuery(); if (JdbcConstants.ORACLE.equals(dbType)) { - return limitOracle(select, dbType, offset, count); + return limitOracle(select, dbType, offset, count, check); } if (JdbcConstants.DB2.equals(dbType)) { - return limitDB2(select, dbType, offset, count); + return limitDB2(select, dbType, offset, count, check); } if (JdbcConstants.SQL_SERVER.equals(dbType) || JdbcUtils.JTDS.equals(dbType)) { - return limitSQLServer(select, dbType, offset, count); + return limitSQLServer(select, dbType, offset, count, check); } if (query instanceof SQLSelectQueryBlock) { - return limitQueryBlock(select, dbType, offset, count); + return limitQueryBlock(select, dbType, offset, count, check); } throw new UnsupportedOperationException(); } - private static String limitQueryBlock(SQLSelect select, String dbType, int offset, int count) { + private static boolean limitQueryBlock(SQLSelect select, String dbType, int offset, int count, boolean check) { SQLSelectQueryBlock queryBlock = (SQLSelectQueryBlock) select.getQuery(); if (JdbcConstants.MYSQL.equals(dbType) || // JdbcConstants.MARIADB.equals(dbType) || // JdbcConstants.H2.equals(dbType)) { - return limitMySqlQueryBlock((MySqlSelectQueryBlock) queryBlock, dbType, offset, count); + return limitMySqlQueryBlock((MySqlSelectQueryBlock) queryBlock, dbType, offset, count, check); } if (JdbcConstants.POSTGRESQL.equals(dbType)) { - return limitPostgreSQLQueryBlock((PGSelectQueryBlock) queryBlock, dbType, offset, count); + return limitPostgreSQLQueryBlock((PGSelectQueryBlock) queryBlock, dbType, offset, count, check); } throw new UnsupportedOperationException(); } - private static String limitPostgreSQLQueryBlock(PGSelectQueryBlock queryBlock, String dbType, int offset, int count) { - if (queryBlock.getLimit() != null) { - throw new IllegalArgumentException("limit already exists."); + private static boolean limitPostgreSQLQueryBlock(PGSelectQueryBlock queryBlock, String dbType, int offset, int count, boolean check) { + SQLLimit limit = queryBlock.getLimit(); + if (limit != null) { + if (offset > 0) { + limit.setOffset(new SQLIntegerExpr(offset)); + } + + if (check && limit.getRowCount() instanceof SQLNumericLiteralExpr) { + int rowCount = ((SQLNumericLiteralExpr) limit.getRowCount()).getNumber().intValue(); + if (rowCount <= count && offset <= 0) { + return false; + } + } + + limit.setRowCount(new SQLIntegerExpr(count)); } - SQLLimit limit = new SQLLimit(); + limit = new SQLLimit(); if (offset > 0) { limit.setOffset(new SQLIntegerExpr(offset)); } limit.setRowCount(new SQLIntegerExpr(count)); queryBlock.setLimit(limit); - - return SQLUtils.toSQLString(queryBlock, dbType); + return true; } - private static String limitDB2(SQLSelect select, String dbType, int offset, int count) { + private static boolean limitDB2(SQLSelect select, String dbType, int offset, int count, boolean check) { SQLSelectQuery query = select.getQuery(); SQLBinaryOpExpr gt = new SQLBinaryOpExpr(new SQLIdentifierExpr("ROWNUM"), // @@ -140,8 +157,15 @@ public class PagerUtils { if (query instanceof SQLSelectQueryBlock) { DB2SelectQueryBlock queryBlock = (DB2SelectQueryBlock) query; if (offset <= 0) { - queryBlock.setFirst(new SQLNumberExpr(count)); - return SQLUtils.toSQLString(select, dbType); + SQLExpr first = queryBlock.getFirst(); + if (check && first != null && first instanceof SQLNumericLiteralExpr) { + int rowCount = ((SQLNumericLiteralExpr) first).getNumber().intValue(); + if (rowCount < count) { + return false; + } + } + queryBlock.setFirst(new SQLIntegerExpr(count)); + return true; } SQLAggregateExpr aggregateExpr = new SQLAggregateExpr("ROW_NUMBER"); @@ -162,11 +186,13 @@ public class PagerUtils { DB2SelectQueryBlock countQueryBlock = new DB2SelectQueryBlock(); countQueryBlock.getSelectList().add(new SQLSelectItem(new SQLAllColumnExpr())); - countQueryBlock.setFrom(new SQLSubqueryTableSource(select, "XX")); + countQueryBlock.setFrom(new SQLSubqueryTableSource(select.clone(), "XX")); countQueryBlock.setWhere(pageCondition); - return SQLUtils.toSQLString(countQueryBlock, dbType); + select.setQuery(countQueryBlock); + + return true; } DB2SelectQueryBlock countQueryBlock = new DB2SelectQueryBlock(); @@ -177,10 +203,11 @@ public class PagerUtils { select.setOrderBy(null); countQueryBlock.getSelectList().add(new SQLSelectItem(aggregateExpr, "ROWNUM")); - countQueryBlock.setFrom(new SQLSubqueryTableSource(select, "XX")); + countQueryBlock.setFrom(new SQLSubqueryTableSource(select.clone(), "XX")); if (offset <= 0) { - return SQLUtils.toSQLString(countQueryBlock, dbType); + select.setQuery(countQueryBlock); + return true; } DB2SelectQueryBlock offsetQueryBlock = new DB2SelectQueryBlock(); @@ -188,10 +215,12 @@ public class PagerUtils { offsetQueryBlock.setFrom(new SQLSubqueryTableSource(new SQLSelect(countQueryBlock), "XXX")); offsetQueryBlock.setWhere(pageCondition); - return SQLUtils.toSQLString(offsetQueryBlock, dbType); + select.setQuery(offsetQueryBlock); + + return true; } - private static String limitSQLServer(SQLSelect select, String dbType, int offset, int count) { + private static boolean limitSQLServer(SQLSelect select, String dbType, int offset, int count, boolean check) { SQLSelectQuery query = select.getQuery(); SQLBinaryOpExpr gt = new SQLBinaryOpExpr(new SQLIdentifierExpr("ROWNUM"), // @@ -208,8 +237,15 @@ public class PagerUtils { if (query instanceof SQLSelectQueryBlock) { SQLServerSelectQueryBlock queryBlock = (SQLServerSelectQueryBlock) query; if (offset <= 0) { + SQLServerTop top = queryBlock.getTop(); + if (check && top != null && !top.isPercent() && top.getExpr() instanceof SQLNumericLiteralExpr) { + int rowCount = ((SQLNumericLiteralExpr) top.getExpr()).getNumber().intValue(); + if (rowCount <= count) { + return false; + } + } queryBlock.setTop(new SQLServerTop(new SQLNumberExpr(count))); - return SQLUtils.toSQLString(select, dbType); + return true; } SQLAggregateExpr aggregateExpr = new SQLAggregateExpr("ROW_NUMBER"); @@ -222,21 +258,25 @@ public class PagerUtils { SQLServerSelectQueryBlock countQueryBlock = new SQLServerSelectQueryBlock(); countQueryBlock.getSelectList().add(new SQLSelectItem(new SQLAllColumnExpr())); - countQueryBlock.setFrom(new SQLSubqueryTableSource(select, "XX")); + countQueryBlock.setFrom(new SQLSubqueryTableSource(select.clone(), "XX")); countQueryBlock.setWhere(pageCondition); - return SQLUtils.toSQLString(countQueryBlock, dbType); + select.setQuery(countQueryBlock); + + return true; } SQLServerSelectQueryBlock countQueryBlock = new SQLServerSelectQueryBlock(); countQueryBlock.getSelectList().add(new SQLSelectItem(new SQLPropertyExpr(new SQLIdentifierExpr("XX"), "*"))); - countQueryBlock.setFrom(new SQLSubqueryTableSource(select, "XX")); + countQueryBlock.setFrom(new SQLSubqueryTableSource(select.clone(), "XX")); if (offset <= 0) { countQueryBlock.setTop(new SQLServerTop(new SQLNumberExpr(count))); - return SQLUtils.toSQLString(countQueryBlock, dbType); + + select.setQuery(countQueryBlock); + return true; } SQLAggregateExpr aggregateExpr = new SQLAggregateExpr("ROW_NUMBER"); @@ -250,15 +290,34 @@ public class PagerUtils { offsetQueryBlock.setFrom(new SQLSubqueryTableSource(new SQLSelect(countQueryBlock), "XXX")); offsetQueryBlock.setWhere(pageCondition); - return SQLUtils.toSQLString(offsetQueryBlock, dbType); + select.setQuery(offsetQueryBlock); + + return true; } - private static String limitOracle(SQLSelect select, String dbType, int offset, int count) { + private static boolean limitOracle(SQLSelect select, String dbType, int offset, int count, boolean check) { SQLSelectQuery query = select.getQuery(); if (query instanceof SQLSelectQueryBlock) { OracleSelectQueryBlock queryBlock = (OracleSelectQueryBlock) query; if (queryBlock.getGroupBy() == null && select.getOrderBy() == null && offset <= 0) { + SQLExpr where = queryBlock.getWhere(); + if (check && where instanceof SQLBinaryOpExpr) { + SQLBinaryOpExpr binaryOpWhere = (SQLBinaryOpExpr) where; + if (binaryOpWhere.getOperator() == SQLBinaryOperator.LessThanOrEqual) { + SQLExpr left = binaryOpWhere.getLeft(); + SQLExpr right = binaryOpWhere.getRight(); + if (left instanceof SQLIdentifierExpr + && ((SQLIdentifierExpr) left).getName().equalsIgnoreCase("ROWNUM") + && right instanceof SQLNumericLiteralExpr) { + int rowCount = ((SQLNumericLiteralExpr) right).getNumber().intValue(); + if (rowCount <= count) { + return false; + } + } + } + } + SQLExpr condition = new SQLBinaryOpExpr(new SQLIdentifierExpr("ROWNUM"), // SQLBinaryOperator.LessThanOrEqual, // new SQLNumberExpr(count), // @@ -272,7 +331,7 @@ public class PagerUtils { JdbcConstants.ORACLE)); } - return SQLUtils.toSQLString(select, dbType); + return true; } } @@ -280,13 +339,16 @@ public class PagerUtils { countQueryBlock.getSelectList().add(new SQLSelectItem(new SQLPropertyExpr(new SQLIdentifierExpr("XX"), "*"))); countQueryBlock.getSelectList().add(new SQLSelectItem(new SQLIdentifierExpr("ROWNUM"), "RN")); - countQueryBlock.setFrom(new SQLSubqueryTableSource(select, "XX")); + countQueryBlock.setFrom(new SQLSubqueryTableSource(select.clone(), "XX")); countQueryBlock.setWhere(new SQLBinaryOpExpr(new SQLIdentifierExpr("ROWNUM"), // SQLBinaryOperator.LessThanOrEqual, // new SQLNumberExpr(count + offset), // JdbcConstants.ORACLE)); + + select.setOrderBy(null); if (offset <= 0) { - return SQLUtils.toSQLString(countQueryBlock, dbType); + select.setQuery(countQueryBlock); + return true; } OracleSelectQueryBlock offsetQueryBlock = new OracleSelectQueryBlock(); @@ -297,22 +359,37 @@ public class PagerUtils { new SQLNumberExpr(offset), // JdbcConstants.ORACLE)); - return SQLUtils.toSQLString(offsetQueryBlock, dbType); + select.setQuery(offsetQueryBlock); + return true; } - private static String limitMySqlQueryBlock(MySqlSelectQueryBlock queryBlock, String dbType, int offset, int count) { - if (queryBlock.getLimit() != null) { - throw new IllegalArgumentException("limit already exists."); + private static boolean limitMySqlQueryBlock(MySqlSelectQueryBlock queryBlock, String dbType, int offset, int count, boolean check) { + SQLLimit limit = queryBlock.getLimit(); + if (limit != null) { + if (offset > 0) { + limit.setOffset(new SQLIntegerExpr(offset)); + } + + if (check && limit.getRowCount() instanceof SQLNumericLiteralExpr) { + int rowCount = ((SQLNumericLiteralExpr) limit.getRowCount()).getNumber().intValue(); + if (rowCount <= count && offset <= 0) { + return false; + } + } + + limit.setRowCount(new SQLIntegerExpr(count)); } - SQLLimit limit = new SQLLimit(); - if (offset > 0) { - limit.setOffset(new SQLNumberExpr(offset)); + if (limit == null) { + limit = new SQLLimit(); + if (offset > 0) { + limit.setOffset(new SQLIntegerExpr(offset)); + } + limit.setRowCount(new SQLIntegerExpr(count)); + queryBlock.setLimit(limit); } - limit.setRowCount(new SQLNumberExpr(count)); - queryBlock.setLimit(limit); - return SQLUtils.toSQLString(queryBlock, dbType); + return true; } private static String count(SQLSelect select, String dbType) { diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java index f475cda5f1..43f7fb9601 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java @@ -133,4 +133,15 @@ public class SQLSelect extends SQLObjectImpl { return super.toString(); } + + public SQLSelect clone() { + SQLSelect clone = new SQLSelect(); + + clone.withSubQuery = this.withSubQuery; + clone.query = this.query; + clone.orderBy = this.orderBy; + clone.hints = this.hints; + + return clone; + } } diff --git a/src/main/java/com/alibaba/druid/wall/WallConfig.java b/src/main/java/com/alibaba/druid/wall/WallConfig.java index 6a10f7cf7f..3ec750becc 100644 --- a/src/main/java/com/alibaba/druid/wall/WallConfig.java +++ b/src/main/java/com/alibaba/druid/wall/WallConfig.java @@ -126,6 +126,8 @@ public class WallConfig implements WallConfigMBean { private boolean completeInsertValuesCheck = false; private int insertValuesCheckSize = 3; + private int selectLimit = -1; + public WallConfig(){ this.configFromProperties(System.getProperties()); } @@ -798,6 +800,14 @@ public class WallConfig implements WallConfigMBean { this.blockAllow = blockAllow; } + public int getSelectLimit() { + return selectLimit; + } + + public void setSelectLimit(int selectLimit) { + this.selectLimit = selectLimit; + } + public void configFromProperties(Properties properties) { { String propertyValue = properties.getProperty("druid.wall.tenantColumn"); diff --git a/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java index 8c3eae096e..b32ada4ef3 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java +++ b/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java @@ -15,6 +15,7 @@ */ package com.alibaba.druid.wall.spi; +import com.alibaba.druid.sql.PagerUtils; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLObject; @@ -23,21 +24,7 @@ import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.expr.SQLInListExpr; import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; -import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCallStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTriggerStatement; -import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; -import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; -import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; -import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; -import com.alibaba.druid.sql.ast.statement.SQLSelectItem; -import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; -import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; -import com.alibaba.druid.sql.ast.statement.SQLSetStatement; -import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; -import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.sql.ast.statement.*; import com.alibaba.druid.sql.dialect.db2.ast.stmt.DB2SelectQueryBlock; import com.alibaba.druid.sql.dialect.db2.visitor.DB2ASTVisitorAdapter; import com.alibaba.druid.util.JdbcConstants; @@ -197,6 +184,13 @@ public class DB2WallVisitor extends DB2ASTVisitorAdapter implements WallVisitor WallVisitorUtils.initWallTopStatementContext(); + int selectLimit = config.getSelectLimit(); + if (selectLimit >= 0) { + SQLSelect select = x.getSelect(); + PagerUtils.limit(select, getDbType(), 0, selectLimit, true); + this.sqlModified = true; + } + return true; } diff --git a/src/main/java/com/alibaba/druid/wall/spi/MySqlWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/MySqlWallVisitor.java index a67e3bee85..c51bc970c0 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/MySqlWallVisitor.java +++ b/src/main/java/com/alibaba/druid/wall/spi/MySqlWallVisitor.java @@ -18,6 +18,7 @@ package com.alibaba.druid.wall.spi; import java.util.ArrayList; import java.util.List; +import com.alibaba.druid.sql.PagerUtils; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLCommentHint; import com.alibaba.druid.sql.ast.SQLName; @@ -28,22 +29,7 @@ import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLNumericLiteralExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr; -import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLAssignItem; -import com.alibaba.druid.sql.ast.statement.SQLCallStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTriggerStatement; -import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; -import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; -import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; -import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; -import com.alibaba.druid.sql.ast.statement.SQLSelectItem; -import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; -import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; -import com.alibaba.druid.sql.ast.statement.SQLSetStatement; -import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; -import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.sql.ast.statement.*; import com.alibaba.druid.sql.dialect.mysql.ast.expr.MySqlOutFileExpr; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlDeleteStatement; @@ -198,6 +184,13 @@ public class MySqlWallVisitor extends MySqlASTVisitorAdapter implements WallVisi } WallVisitorUtils.initWallTopStatementContext(); + + int selectLimit = config.getSelectLimit(); + if (selectLimit >= 0) { + SQLSelect select = x.getSelect(); + PagerUtils.limit(select, getDbType(), 0, selectLimit, true); + this.sqlModified = true; + } return true; } diff --git a/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java index 0d7ff32954..72f1d8a93f 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java +++ b/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java @@ -18,6 +18,7 @@ package com.alibaba.druid.wall.spi; import java.util.ArrayList; import java.util.List; +import com.alibaba.druid.sql.PagerUtils; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLObject; @@ -26,21 +27,7 @@ import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.expr.SQLInListExpr; import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; -import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCallStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTriggerStatement; -import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; -import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; -import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; -import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; -import com.alibaba.druid.sql.ast.statement.SQLSelectItem; -import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; -import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; -import com.alibaba.druid.sql.ast.statement.SQLSetStatement; -import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; -import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.sql.ast.statement.*; import com.alibaba.druid.sql.dialect.oracle.ast.stmt.OracleCreateTableStatement; import com.alibaba.druid.sql.dialect.oracle.ast.stmt.OracleDeleteStatement; import com.alibaba.druid.sql.dialect.oracle.ast.stmt.OracleInsertStatement; @@ -208,6 +195,13 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi } WallVisitorUtils.initWallTopStatementContext(); + int selectLimit = config.getSelectLimit(); + if (selectLimit >= 0) { + SQLSelect select = x.getSelect(); + PagerUtils.limit(select, getDbType(), 0, selectLimit, true); + this.sqlModified = true; + } + return true; } diff --git a/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java index 9b253751d7..e46b0c2b01 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java +++ b/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java @@ -15,6 +15,7 @@ */ package com.alibaba.druid.wall.spi; +import com.alibaba.druid.sql.PagerUtils; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLObject; @@ -23,22 +24,9 @@ import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.expr.SQLInListExpr; import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; -import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCallStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTriggerStatement; -import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; -import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; -import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; -import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; -import com.alibaba.druid.sql.ast.statement.SQLSelectItem; -import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; -import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; -import com.alibaba.druid.sql.ast.statement.SQLSetStatement; -import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; -import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.sql.ast.statement.*; import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock; +import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement; import com.alibaba.druid.sql.dialect.postgresql.visitor.PGASTVisitorAdapter; import com.alibaba.druid.util.JdbcConstants; import com.alibaba.druid.wall.Violation; @@ -194,6 +182,33 @@ public class PGWallVisitor extends PGASTVisitorAdapter implements WallVisitor { WallVisitorUtils.initWallTopStatementContext(); + int selectLimit = config.getSelectLimit(); + if (selectLimit >= 0) { + SQLSelect select = x.getSelect(); + PagerUtils.limit(select, getDbType(), 0, selectLimit); + this.sqlModified = true; + } + + return true; + } + + @Override + public boolean visit(PGSelectStatement x) { + if (!config.isSelelctAllow()) { + this.getViolations().add(new IllegalSQLObjectViolation(ErrorCode.SELECT_NOT_ALLOW, "selelct not allow", + this.toSQL(x))); + return false; + } + + WallVisitorUtils.initWallTopStatementContext(); + + int selectLimit = config.getSelectLimit(); + if (selectLimit >= 0) { + SQLSelect select = x.getSelect(); + PagerUtils.limit(select, getDbType(), 0, selectLimit, true); + this.sqlModified = true; + } + return true; } diff --git a/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java index 0c959c8a44..c1539b85c9 100644 --- a/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java +++ b/src/main/java/com/alibaba/druid/wall/spi/SQLServerWallVisitor.java @@ -15,6 +15,7 @@ */ package com.alibaba.druid.wall.spi; +import com.alibaba.druid.sql.PagerUtils; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLObject; @@ -24,21 +25,7 @@ import com.alibaba.druid.sql.ast.expr.SQLInListExpr; import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr; -import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCallStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLCreateTriggerStatement; -import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; -import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement; -import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; -import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; -import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; -import com.alibaba.druid.sql.ast.statement.SQLSelectItem; -import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; -import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; -import com.alibaba.druid.sql.ast.statement.SQLSetStatement; -import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; -import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.sql.ast.statement.*; import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock; import com.alibaba.druid.sql.dialect.sqlserver.ast.expr.SQLServerObjectReferenceExpr; import com.alibaba.druid.sql.dialect.sqlserver.ast.stmt.SQLServerExecStatement; @@ -206,6 +193,13 @@ public class SQLServerWallVisitor extends SQLServerASTVisitorAdapter implements } WallVisitorUtils.initWallTopStatementContext(); + int selectLimit = config.getSelectLimit(); + if (selectLimit >= 0) { + SQLSelect select = x.getSelect(); + PagerUtils.limit(select, getDbType(), 0, selectLimit, true); + this.sqlModified = true; + } + return true; } diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest.java new file mode 100644 index 0000000000..be4e7ae1b3 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest.java @@ -0,0 +1,83 @@ +/* + * Copyright 1999-2017 Alibaba Group Holding Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.druid.bvt.filter.wall; + +import com.alibaba.druid.wall.WallCheckResult; +import com.alibaba.druid.wall.WallConfig; +import com.alibaba.druid.wall.WallProvider; +import com.alibaba.druid.wall.WallUtils; +import com.alibaba.druid.wall.spi.*; +import junit.framework.TestCase; +import org.junit.Assert; + +public class WallSelectLimitTest extends TestCase { + private String sql = "select * from t"; + + private WallConfig config = new WallConfig(); + + protected void setUp() throws Exception { + config.setSelectLimit(1000); + } + + public void testMySql() throws Exception { + WallProvider provider = new MySqlWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "LIMIT 1000", resultSql); + } + + public void testDB2() throws Exception { + WallProvider provider = new DB2WallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "FETCH FIRST 1000 ROWS ONLY", resultSql); + } + + public void testSQLServer() throws Exception { + WallProvider provider = new SQLServerWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT TOP 1000 *\n" + + "FROM t", resultSql); + } + + public void testOracle() throws Exception { + WallProvider provider = new OracleWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "WHERE ROWNUM <= 1000", resultSql); + } + + public void testPG() throws Exception { + WallProvider provider = new PGWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "LIMIT 1000", resultSql); + } +} diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest_2.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest_2.java new file mode 100644 index 0000000000..e9bbdb5e2d --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/WallSelectLimitTest_2.java @@ -0,0 +1,87 @@ +/* + * Copyright 1999-2017 Alibaba Group Holding Ltd. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.druid.bvt.filter.wall; + +import com.alibaba.druid.sql.PagerUtils; +import com.alibaba.druid.util.JdbcConstants; +import com.alibaba.druid.wall.WallCheckResult; +import com.alibaba.druid.wall.WallConfig; +import com.alibaba.druid.wall.WallProvider; +import com.alibaba.druid.wall.spi.*; +import junit.framework.TestCase; +import sun.jvm.hotspot.debugger.Page; + +public class WallSelectLimitTest_2 extends TestCase { + private WallConfig config = new WallConfig(); + + protected void setUp() throws Exception { + config.setSelectLimit(1000); + } + + public void testMySql() throws Exception { + String sql = "select * from t limit 10"; + WallProvider provider = new MySqlWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "LIMIT 10", resultSql); + } + + public void testPG() throws Exception { + String sql = "select * from t limit 10"; + WallProvider provider = new PGWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "LIMIT 10", resultSql); + } + + public void testDB2() throws Exception { + String sql = PagerUtils.limit("select * from t", JdbcConstants.DB2, 0, 10); + WallProvider provider = new DB2WallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "FETCH FIRST 10 ROWS ONLY", resultSql); + } + + public void testSQLServer() throws Exception { + String sql = PagerUtils.limit("select * from t", JdbcConstants.SQL_SERVER, 0, 10); + WallProvider provider = new SQLServerWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT TOP 10 *\n" + + "FROM t", resultSql); + } + + public void testOracle() throws Exception { + String sql = PagerUtils.limit("select * from t", JdbcConstants.ORACLE, 0, 10); + WallProvider provider = new OracleWallProvider(config); + WallCheckResult checkResult = provider.check(sql); + String resultSql = checkResult.getSql(); + System.out.println(resultSql); + assertEquals("SELECT *\n" + + "FROM t\n" + + "WHERE ROWNUM <= 10", resultSql); + } +}