From 46c9affa838b779b2623b4b5dc7086c27ec99e1b Mon Sep 17 00:00:00 2001
From: Amos Bird <amosbird@gmail.com>
Date: Mon, 5 Apr 2021 14:36:33 +0800
Subject: [PATCH] Fix CTE in view.

---
 src/Interpreters/ApplyWithSubqueryVisitor.cpp | 53 +++++++++++--------
 src/Interpreters/ApplyWithSubqueryVisitor.h   |  6 +++
 src/Interpreters/InterpreterCreateQuery.cpp   |  3 ++
 .../0_stateless/01711_cte_subquery_fix.sql    |  4 ++
 4 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.cpp b/src/Interpreters/ApplyWithSubqueryVisitor.cpp
index f3d0c82e5f1..27e0e438a71 100644
--- a/src/Interpreters/ApplyWithSubqueryVisitor.cpp
+++ b/src/Interpreters/ApplyWithSubqueryVisitor.cpp
@@ -13,28 +13,7 @@ namespace DB
 void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
 {
     if (auto * node_select = ast->as<ASTSelectQuery>())
-    {
-        std::optional<Data> new_data;
-        if (auto with = node_select->with())
-        {
-            for (auto & child : with->children)
-            {
-                visit(child, new_data ? *new_data : data);
-                if (auto * ast_with_elem = child->as<ASTWithElement>())
-                {
-                    if (!new_data)
-                        new_data = data;
-                    new_data->subqueries[ast_with_elem->name] = ast_with_elem->subquery;
-                }
-            }
-        }
-
-        for (auto & child : node_select->children)
-        {
-            if (child != node_select->with())
-                visit(child, new_data ? *new_data : data);
-        }
-    }
+        visit(*node_select, data);
     else
     {
         for (auto & child : ast->children)
@@ -46,6 +25,36 @@ void ApplyWithSubqueryVisitor::visit(ASTPtr & ast, const Data & data)
     }
 }
 
+void ApplyWithSubqueryVisitor::visit(ASTSelectQuery & ast, const Data & data)
+{
+    std::optional<Data> new_data;
+    if (auto with = ast.with())
+    {
+        for (auto & child : with->children)
+        {
+            visit(child, new_data ? *new_data : data);
+            if (auto * ast_with_elem = child->as<ASTWithElement>())
+            {
+                if (!new_data)
+                    new_data = data;
+                new_data->subqueries[ast_with_elem->name] = ast_with_elem->subquery;
+            }
+        }
+    }
+
+    for (auto & child : ast.children)
+    {
+        if (child != ast.with())
+            visit(child, new_data ? *new_data : data);
+    }
+}
+
+void ApplyWithSubqueryVisitor::visit(ASTSelectWithUnionQuery & ast, const Data & data)
+{
+    for (auto & child : ast.children)
+        visit(child, data);
+}
+
 void ApplyWithSubqueryVisitor::visit(ASTTableExpression & table, const Data & data)
 {
     if (table.database_and_table_name)
diff --git a/src/Interpreters/ApplyWithSubqueryVisitor.h b/src/Interpreters/ApplyWithSubqueryVisitor.h
index ed3f9fbcd3e..3b17de1d691 100644
--- a/src/Interpreters/ApplyWithSubqueryVisitor.h
+++ b/src/Interpreters/ApplyWithSubqueryVisitor.h
@@ -7,6 +7,8 @@
 namespace DB
 {
 class ASTFunction;
+class ASTSelectQuery;
+class ASTSelectWithUnionQuery;
 struct ASTTableExpression;
 
 class ApplyWithSubqueryVisitor
@@ -18,9 +20,13 @@ public:
     };
 
     static void visit(ASTPtr & ast) { visit(ast, {}); }
+    static void visit(ASTSelectQuery & select) { visit(select, {}); }
+    static void visit(ASTSelectWithUnionQuery & select) { visit(select, {}); }
 
 private:
     static void visit(ASTPtr & ast, const Data & data);
+    static void visit(ASTSelectQuery & ast, const Data & data);
+    static void visit(ASTSelectWithUnionQuery & ast, const Data & data);
     static void visit(ASTTableExpression & table, const Data & data);
     static void visit(ASTFunction & func, const Data & data);
 };
diff --git a/src/Interpreters/InterpreterCreateQuery.cpp b/src/Interpreters/InterpreterCreateQuery.cpp
index 8cd93546ad0..26ab435aaf1 100644
--- a/src/Interpreters/InterpreterCreateQuery.cpp
+++ b/src/Interpreters/InterpreterCreateQuery.cpp
@@ -59,6 +59,7 @@
 #include <Interpreters/QueryLog.h>
 #include <Interpreters/addTypeConversionToAST.h>
 #include <Interpreters/FunctionNameNormalizer.h>
+#include <Interpreters/ApplyWithSubqueryVisitor.h>
 
 #include <TableFunctions/TableFunctionFactory.h>
 #include <common/logger_useful.h>
@@ -887,6 +888,8 @@ BlockIO InterpreterCreateQuery::createTable(ASTCreateQuery & create)
 
     if (create.select && create.isView())
     {
+        // Expand CTE before filling default database
+        ApplyWithSubqueryVisitor().visit(*create.select);
         AddDefaultDatabaseVisitor visitor(current_database);
         visitor.visit(*create.select);
     }
diff --git a/tests/queries/0_stateless/01711_cte_subquery_fix.sql b/tests/queries/0_stateless/01711_cte_subquery_fix.sql
index ddea548eada..10ad9019209 100644
--- a/tests/queries/0_stateless/01711_cte_subquery_fix.sql
+++ b/tests/queries/0_stateless/01711_cte_subquery_fix.sql
@@ -1,3 +1,7 @@
 drop table if exists t;
 create table t engine = Memory as with cte as (select * from numbers(10)) select * from cte;
 drop table t;
+
+drop table if exists view1;
+create view view1 as with t as (select number n from numbers(3)) select n from t;
+drop table view1;