Merge pull request #262 from wenshao/master

add PagerUtils

src/main/java/com/alibaba/druid/sql/PagerUtils.java

@@ -0,0 +1,134 @@
+package com.alibaba.druid.sql;
+
+import java.util.List;
+
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr;
+import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr;
+import com.alibaba.druid.sql.ast.statement.SQLSelect;
+import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
+import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
+import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
+import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
+import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLUnionQuery;
+import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
+import com.alibaba.druid.sql.dialect.oracle.ast.expr.OracleAggregateExpr;
+import com.alibaba.druid.sql.dialect.oracle.ast.stmt.OracleSelectQueryBlock;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock;
+import com.alibaba.druid.sql.dialect.sqlserver.ast.SQLServerSelectQueryBlock;
+import com.alibaba.druid.util.JdbcConstants;
+
+public class PagerUtils {
+
+    public static String count(String sql, String dbType) {
+        List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, dbType);
+
+        if (stmtList.size() != 1) {
+            throw new IllegalArgumentException("sql not support count : " + sql);
+        }
+
+        SQLStatement stmt = stmtList.get(0);
+
+        if (!(stmt instanceof SQLSelectStatement)) {
+            throw new IllegalArgumentException("sql not support count : " + sql);
+        }
+
+        SQLSelectStatement selectStmt = (SQLSelectStatement) stmt;
+
+        return count(selectStmt, dbType);
+    }
+
+    private static String count(SQLSelectStatement stmt, String dbType) {
+        SQLSelect select = stmt.getSelect();
+        if (select.getOrderBy() != null) {
+            select.setOrderBy(null);
+        }
+        SQLSelectQuery query = select.getQuery();
+        clearOrderBy(query);
+
+        if (query instanceof SQLSelectQueryBlock) {
+            SQLSelectItem countItem = createCountItem(dbType);
+            
+            SQLSelectQueryBlock queryBlock = (SQLSelectQueryBlock) query;
+            queryBlock.getSelectList().clear();
+            queryBlock.getSelectList().add(countItem);
+            return SQLUtils.toSQLString(stmt, dbType);
+        } else if (query instanceof SQLUnionQuery) {
+            SQLSelectQueryBlock countSelectQuery = createQueryBlock(dbType);
+            
+            SQLSelectItem countItem = createCountItem(dbType);
+            countSelectQuery.getSelectList().add(countItem);
+            
+            countSelectQuery.setFrom(new SQLSubqueryTableSource(select));
+            
+            SQLSelect countSelect = new SQLSelect(countSelectQuery);
+            SQLSelectStatement countStmt = new SQLSelectStatement(countSelect);
+            
+            return SQLUtils.toSQLString(countStmt, dbType);
+        }
+
+        throw new IllegalStateException();
+    }
+
+    private static SQLSelectQueryBlock createQueryBlock(String dbType) {
+        if (JdbcConstants.MYSQL.equals(dbType)) {
+            return new MySqlSelectQueryBlock();
+        }
+        
+        if (JdbcConstants.ORACLE.equals(dbType)) {
+            return new OracleSelectQueryBlock();
+        }
+        
+        if (JdbcConstants.POSTGRESQL.equals(dbType)) {
+            return new PGSelectQueryBlock();
+        }
+        
+        if (JdbcConstants.SQL_SERVER.equals(dbType)) {
+            return new SQLServerSelectQueryBlock();
+        }
+        
+        return new SQLSelectQueryBlock();
+    }
+
+    private static SQLSelectItem createCountItem(String dbType) {
+        SQLAggregateExpr countExpr;
+
+        if (JdbcConstants.ORACLE.equals(dbType)) {
+            countExpr = new OracleAggregateExpr("COUNT");
+        } else {
+            countExpr = new SQLAggregateExpr("COUNT");
+        }
+        countExpr.getArguments().add(new SQLAllColumnExpr());
+
+        SQLSelectItem countItem = new SQLSelectItem(countExpr);
+        return countItem;
+    }
+
+    private static void clearOrderBy(SQLSelectQuery query) {
+        if (query instanceof SQLSelectQueryBlock) {
+            SQLSelectQueryBlock queryBlock = (SQLSelectQueryBlock) query;
+            if (queryBlock instanceof MySqlSelectQueryBlock) {
+                MySqlSelectQueryBlock mysqlQueryBlock = (MySqlSelectQueryBlock) queryBlock;
+                if (mysqlQueryBlock.getOrderBy() != null) {
+                    mysqlQueryBlock.setOrderBy(null);
+                }
+            } else if (queryBlock instanceof PGSelectQueryBlock) {
+                PGSelectQueryBlock pgQueryBlock = (PGSelectQueryBlock) queryBlock;
+                if (pgQueryBlock.getOrderBy() != null) {
+                    pgQueryBlock.setOrderBy(null);
+                }
+            }
+            return;
+        }
+
+        if (query instanceof SQLUnionQuery) {
+            SQLUnionQuery union = (SQLUnionQuery) query;
+            if (union.getOrderBy() != null) {
+                union.setOrderBy(null);
+            }
+            clearOrderBy(union.getLeft());
+            clearOrderBy(union.getRight());
+        }
+    }
+}

src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java

@@ -30,6 +30,10 @@ public class SQLSelect extends SQLObjectImpl {
     public SQLSelect(){
 
     }
+    
+    public SQLSelect(SQLSelectQuery query) {
+        this.query = query;
+    }
 
     public SQLWithSubqueryClause getWithSubQuery() {
         return withSubQuery;

src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlCreateTableParser.java

@@ -180,7 +180,6 @@ public class MySqlCreateTableParser extends SQLCreateTableParser {
         return stmt;
     }
 
-    @SuppressWarnings("unused")
     protected SQLTableConstaint parseConstraint() {
         SQLName name = null;
         if (lexer.token() == (Token.CONSTRAINT)) {

src/test/java/PagerUtils/PagerUtilsTest0.java

@@ -0,0 +1,44 @@
+package PagerUtils;
+
+import junit.framework.TestCase;
+
+import org.junit.Assert;
+
+import com.alibaba.druid.sql.PagerUtils;
+import com.alibaba.druid.util.JdbcConstants;
+
+public class PagerUtilsTest0 extends TestCase {
+
+    public void test_mysql_0() throws Exception {
+        String sql = "select * from t";
+        String result = PagerUtils.count(sql, JdbcConstants.MYSQL);
+        Assert.assertEquals("SELECT COUNT(*)\n" + //
+                            "FROM t", result);
+    }
+
+    public void test_mysql_1() throws Exception {
+        String sql = "select id, name from t";
+        String result = PagerUtils.count(sql, JdbcConstants.MYSQL);
+        Assert.assertEquals("SELECT COUNT(*)\n" + //
+                            "FROM t", result);
+    }
+
+    public void test_mysql_2() throws Exception {
+        String sql = "select id, name from t order by id";
+        String result = PagerUtils.count(sql, JdbcConstants.MYSQL);
+        Assert.assertEquals("SELECT COUNT(*)\n" + //
+                            "FROM t", result);
+    }
+
+    public void test_mysql_3() throws Exception {
+        String sql = "select id, name from t1 union select id, name from t2 order by id";
+        String result = PagerUtils.count(sql, JdbcConstants.MYSQL);
+        Assert.assertEquals("SELECT COUNT(*)\n" + //
+                            "FROM (SELECT id, name" + //
+                            "\n\tFROM t1" + //
+                            "\n\tUNION" + //
+                            "\n\tSELECT id, name" + //
+                            "\n\tFROM t2" + //
+                            "\n\t)", result);
+    }
+}