Merge

dbms/include/DB/Parsers/ASTQueryWithOutput.h

@@ -5,10 +5,9 @@
 
 namespace DB
 {
-	
-	
+
 /** Запрос с секцией FORMAT.
-	*/
+  */
 class ASTQueryWithOutput : public IAST
 {
 public:
@@ -16,6 +15,11 @@ public:
 
 	ASTQueryWithOutput() = default;
 	ASTQueryWithOutput(const StringRange range_) : IAST(range_) {}
+
+	/** Возвращает указатель на формат. Если типом объекта является ASTSelectQuery,
+	  * то эта функция возвращает указатель на формат из последнего SELECT'а цепочки UNION ALL.
+	  */
+	virtual const IAST * getFormat() const { return format.get(); }
 };
 
 

dbms/include/DB/Parsers/ASTSelectQuery.h

@@ -34,6 +34,9 @@ public:
 
 	ASTPtr clone() const override;
 
+	/// Возвращает указатель на формат из последнего SELECT'а цепочки UNION ALL.
+	const IAST * getFormat() const override;
+
 public:
 	bool distinct = false;
 	ASTPtr select_expression_list;

dbms/src/Client/Client.cpp

@@ -850,11 +850,11 @@ private:
 			/// Формат может быть указан в запросе.
 			if (ASTQueryWithOutput * query_with_output = dynamic_cast<ASTQueryWithOutput *>(&*parsed_query))
 			{
-				if (query_with_output->format)
+				if (query_with_output->getFormat() != nullptr)
 				{
 					if (has_vertical_output_suffix)
 						throw Exception("Output format already specified", ErrorCodes::CLIENT_OUTPUT_FORMAT_SPECIFIED);
-					if (ASTIdentifier * id = typeid_cast<ASTIdentifier *>(&*query_with_output->format))
+					if (const ASTIdentifier * id = typeid_cast<const ASTIdentifier *>(query_with_output->getFormat()))
 						current_format = id->name;
 				}
 			}

dbms/src/Interpreters/executeQuery.cpp

@@ -211,8 +211,8 @@ void executeQuery(
 	{
 		const ASTQueryWithOutput * ast_query_with_output = dynamic_cast<const ASTQueryWithOutput *>(ast.get());
 
-		String format_name = ast_query_with_output && ast_query_with_output->format
-			? typeid_cast<const ASTIdentifier &>(*ast_query_with_output->format).name
+		String format_name = ast_query_with_output && (ast_query_with_output->getFormat() != nullptr)
+			? typeid_cast<const ASTIdentifier &>(*ast_query_with_output->getFormat()).name
 			: context.getDefaultFormat();
 
 		BlockOutputStreamPtr out = context.getFormatFactory().getOutput(format_name, ostr, streams.in_sample);

dbms/src/Parsers/ASTSelectQuery.cpp

@@ -172,5 +172,14 @@ ASTPtr ASTSelectQuery::clone() const
 
 	return ptr;
 }
+
+const IAST * ASTSelectQuery::getFormat() const
+{
+	const ASTSelectQuery * query = this;
+	while (!query->next_union_all.isNull())
+		query = static_cast<const ASTSelectQuery *>(query->next_union_all.get());
+	return query->format.get();
+}
+
 };
 

dbms/src/Parsers/ParserSelectQuery.cpp

@@ -8,7 +8,6 @@
 #include <DB/Parsers/ParserSetQuery.h>
 #include <DB/Parsers/ParserSelectQuery.h>
 
-
 namespace DB
 {
 
@@ -296,6 +295,8 @@ bool ParserSelectQuery::parseImpl(Pos & pos, Pos end, ASTPtr & node, Pos & max_p
 		ws.ignore(pos, end);
 	}
 
+	bool has_format = false;
+
 	/// FORMAT format_name
 	if (s_format.ignore(pos, end, max_parsed_pos, expected))
 	{
@@ -308,6 +309,7 @@ bool ParserSelectQuery::parseImpl(Pos & pos, Pos end, ASTPtr & node, Pos & max_p
 		typeid_cast<ASTIdentifier &>(*select_query->format).kind = ASTIdentifier::Format;
 
 		ws.ignore(pos, end);
+		has_format = true;
 	}
 
 	// UNION ALL select query
@@ -317,6 +319,13 @@ bool ParserSelectQuery::parseImpl(Pos & pos, Pos end, ASTPtr & node, Pos & max_p
 
 		if (s_all.ignore(pos, end, max_parsed_pos, expected))
 		{
+			if (has_format)
+			{
+				/// FORMAT может быть задан только в последнем запросе цепочки UNION ALL.
+				expected = "FORMAT only in the last SELECT of the UNION ALL chain";
+				return false;
+			}
+
 			ParserSelectQuery select_p;
 			if (!select_p.parse(pos, end, select_query->next_union_all, max_parsed_pos, expected))
 				return false;

dbms/tests/queries/0_stateless/00098_k_union_all.reference

@@ -0,0 +1,18 @@
+1
+
+1
+1
+
+1
+1
+
+2
+1
+
+1
+1
+
+2
+1
+
+3

dbms/tests/queries/0_stateless/00098_k_union_all.sql

@@ -0,0 +1,3 @@
+SELECT 1 FORMAT PrettySpace;
+SELECT 1 UNION ALL SELECT 2 FORMAT PrettySpace;
+SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 FORMAT PrettySpace;