From fcc6fcaa9b4ae8836d305754d2580222dfb8e245 Mon Sep 17 00:00:00 2001
From: wenshao <szujobs@hotmail.com>
Date: Wed, 1 Oct 2014 04:46:42 +0800
Subject: [PATCH] bug fixed for SchemaStat

---
 .../druid/sql/parser/SQLSelectParser.java     |  1 +
 .../druid/sql/visitor/SchemaStatVisitor.java  | 27 +++++++++++++++++++
 .../bvt/sql/postgresql/PGSelectTest5.java     | 13 +++++++--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java
index ecb0894012..475900cf40 100644
--- a/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java
+++ b/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java
@@ -202,6 +202,7 @@ public class SQLSelectParser extends SQLParser {
 
             for (;;) {
                 SQLWithSubqueryClause.Entry entry = new SQLWithSubqueryClause.Entry();
+                entry.setParent(withQueryClause);
                 entry.setName((SQLIdentifierExpr) this.exprParser.name());
 
                 if (lexer.token() == Token.LPAREN) {
diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java
index 6f24538a2c..99d6b4f943 100644
--- a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java
+++ b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java
@@ -83,6 +83,7 @@ import com.alibaba.druid.sql.ast.statement.SQLTableElement;
 import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement;
 import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
 import com.alibaba.druid.sql.ast.statement.SQLUseStatement;
+import com.alibaba.druid.sql.ast.statement.SQLWithSubqueryClause;
 import com.alibaba.druid.stat.TableStat;
 import com.alibaba.druid.stat.TableStat.Column;
 import com.alibaba.druid.stat.TableStat.Condition;
@@ -872,6 +873,32 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter {
 
     public void endVisit(SQLSelectStatement x) {
     }
+    
+    @Override
+    public boolean visit(SQLWithSubqueryClause.Entry x) {
+        String alias = x.getName().toString();
+        Map<String, String> aliasMap = getAliasMap();
+        SQLWithSubqueryClause with = (SQLWithSubqueryClause) x.getParent();
+
+        if (Boolean.TRUE == with.getRecursive()) {
+            
+            if (aliasMap != null && alias != null) {
+                aliasMap.put(alias, null);
+                subQueryMap.put(alias, x.getSubQuery().getQuery());
+            }
+            
+            x.getSubQuery().accept(this);
+        } else {
+            x.getSubQuery().accept(this);
+            
+            if (aliasMap != null && alias != null) {
+                aliasMap.put(alias, null);
+                subQueryMap.put(alias, x.getSubQuery().getQuery());
+            }
+        }
+        
+        return false;
+    }
 
     public boolean visit(SQLSubqueryTableSource x) {
         x.getSelect().accept(this);
diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectTest5.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectTest5.java
index f2824ae1b0..68faaefc5a 100644
--- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectTest5.java
+++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectTest5.java
@@ -27,7 +27,16 @@ import com.alibaba.druid.sql.dialect.postgresql.visitor.PGSchemaStatVisitor;
 public class PGSelectTest5 extends PGTest {
 
     public void test_0() throws Exception {
-        String sql = " SELECT rs.* FROM ( WITH RECURSIVE r AS( SELECT * FROM t_e_shopcatalog WHERE parentcatalogid= (SELECT catalogid FROM t_e_shopcatalog where catalogname='学习分类' and parentcatalogid='0') UNION ALL SELECT t_e_shopcatalog.* FROM t_e_shopcatalog, r WHERE t_e_shopcatalog.parentcatalogid = r.catalogid )SELECT * FROM r )rs WHERE 1=1";
+        String sql = " SELECT rs.* FROM ( " //
+                     + "WITH RECURSIVE r AS( "//
+                     + "SELECT * " //
+                     + "FROM t_e_shopcatalog " //
+                     + "WHERE parentcatalogid= (SELECT catalogid FROM t_e_shopcatalog where catalogname='学习分类' and parentcatalogid='0') " //
+                     + "UNION ALL " //
+                     + "SELECT t_e_shopcatalog.* " //
+                     + "FROM t_e_shopcatalog, r " //
+                     + "WHERE t_e_shopcatalog.parentcatalogid = r.catalogid )"
+                     + "SELECT * FROM r )rs WHERE 1=1";
 
         PGSQLStatementParser parser = new PGSQLStatementParser(sql);
         List<SQLStatement> statementList = parser.parseStatementList();
@@ -43,7 +52,7 @@ public class PGSelectTest5 extends PGTest {
         System.out.println("fields : " + visitor.getColumns());
         System.out.println("coditions : " + visitor.getConditions());
 
-        Assert.assertEquals(1, visitor.getColumns().size());
+        Assert.assertEquals(4, visitor.getColumns().size());
         Assert.assertEquals(1, visitor.getTables().size());
     }