ISSUES-2133 Support union query with subquery

dbms/src/Parsers/ParserSelectWithUnionQuery.cpp

@@ -1,7 +1,9 @@
-#include <Parsers/ParserSelectQuery.h>
 #include <Parsers/ExpressionListParsers.h>
 #include <Parsers/ParserSelectWithUnionQuery.h>
 #include <Parsers/ASTSelectWithUnionQuery.h>
+#include <Parsers/ParserUnionQueryElement.h>
+#include <Parsers/ASTExpressionList.h>
+#include <Common/typeid_cast.h>
 
 
 namespace DB
@@ -11,17 +13,34 @@ bool ParserSelectWithUnionQuery::parseImpl(Pos & pos, ASTPtr & node, Expected &
 {
     ASTPtr list_node;
 
-    ParserList parser(std::make_unique<ParserSelectQuery>(), std::make_unique<ParserKeyword>("UNION ALL"), false);
+    ParserList parser(std::make_unique<ParserUnionQueryElement>(), std::make_unique<ParserKeyword>("UNION ALL"), false);
     if (!parser.parse(pos, list_node, expected))
         return false;
 
-    auto res = std::make_shared<ASTSelectWithUnionQuery>();
+    auto select_with_union_query = std::make_shared<ASTSelectWithUnionQuery>();
 
-    res->list_of_selects = std::move(list_node);
+    node = select_with_union_query;
-    res->children.push_back(res->list_of_selects);
+    select_with_union_query->list_of_selects = std::make_shared<ASTExpressionList>();
+    select_with_union_query->children.push_back(select_with_union_query->list_of_selects);
+
+    // flatten inner union query
+    for (auto & child : list_node->children)
+        getSelectsFromUnionListNode(child, select_with_union_query->list_of_selects->children);
 
-    node = res;
     return true;
 }
 
+void ParserSelectWithUnionQuery::getSelectsFromUnionListNode(ASTPtr & ast_select, ASTs & selects)
+{
+    if (ASTSelectWithUnionQuery * inner_union = typeid_cast<ASTSelectWithUnionQuery *>(ast_select.get()))
+    {
+        for (auto & child : inner_union->list_of_selects->children)
+            getSelectsFromUnionListNode(child, selects);
+
+        return;
+    }
+
+    selects.push_back(std::move(ast_select));
+}
+
 }

dbms/src/Parsers/ParserSelectWithUnionQuery.h

@@ -1,7 +1,6 @@
 #pragma once
 
-#include <Parsers/ParserQueryWithOutput.h>
+#include <Parsers/IParserBase.h>
-
 
 namespace DB
 {
@@ -12,6 +11,9 @@ class ParserSelectWithUnionQuery : public IParserBase
 protected:
     const char * getName() const override { return "SELECT query, possibly with UNION"; }
     bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected) override;
+
+private:
+    void getSelectsFromUnionListNode(ASTPtr & ast_select, ASTs & selects);
 };
 
 }

dbms/src/Parsers/ParserUnionQueryElement.cpp

@@ -0,0 +1,22 @@
+#include <Parsers/ASTSubquery.h>
+#include <Parsers/ExpressionElementParsers.h>
+#include <Parsers/ParserSelectQuery.h>
+#include <Parsers/ParserUnionQueryElement.h>
+#include <Common/typeid_cast.h>
+
+
+namespace DB
+{
+
+bool ParserUnionQueryElement::parseImpl(Pos & pos, ASTPtr & node, Expected & expected)
+{
+    if (!ParserSubquery().parse(pos, node, expected) && !ParserSelectQuery().parse(pos, node, expected))
+        return false;
+
+    if (auto * ast_sub_query = typeid_cast<ASTSubquery *>(node.get()))
+        node = ast_sub_query->children.at(0);
+
+    return true;
+}
+
+}

dbms/src/Parsers/ParserUnionQueryElement.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <Parsers/IParserBase.h>
+
+
+namespace DB
+{
+
+
+class ParserUnionQueryElement : public IParserBase
+{
+protected:
+    const char * getName() const override { return "SELECT query, subquery, possibly with UNION"; }
+    bool parseImpl(Pos & pos, ASTPtr & node, Expected & expected) override;
+};
+
+}

dbms/tests/queries/0_stateless/00612_union_query_with_subquery.reference

@@ -0,0 +1,12 @@
+0
+0
+0
+1
+1
+2
+0
+0
+0
+1
+1
+2

dbms/tests/queries/0_stateless/00612_union_query_with_subquery.sql

@@ -0,0 +1,2 @@
+SELECT * FROM ((SELECT * FROM system.numbers LIMIT 1) UNION ALL SELECT * FROM system.numbers LIMIT 2 UNION ALL (SELECT * FROM system.numbers LIMIT 3)) ORDER BY number;
+SELECT * FROM (SELECT * FROM system.numbers LIMIT 1 UNION ALL (SELECT * FROM system.numbers LIMIT 2 UNION ALL (SELECT * FROM system.numbers LIMIT 3))) ORDER BY number;