Parse compound select statements

moor_generator/lib/src/analyzer/sql_queries/query_handler.dart

@@ -20,7 +20,7 @@ class QueryHandler {
   Iterable<FoundVariable> get _foundVariables =>
       _foundElements.whereType<FoundVariable>();
 
-  SelectStatement get _select => context.root as SelectStatement;
+  BaseSelectStatement get _select => context.root as BaseSelectStatement;
 
   QueryHandler(this.name, this.context, this.mapper);
 
@@ -39,7 +39,7 @@ class QueryHandler {
 
   SqlQuery _mapToMoor() {
     final root = context.root;
-    if (root is SelectStatement) {
+    if (root is BaseSelectStatement) {
       return _handleSelect();
     } else if (root is UpdateStatement ||
         root is DeleteStatement ||

sqlparser/lib/src/ast/ast.dart

@@ -152,6 +152,8 @@ abstract class AstNode {
 
 abstract class AstVisitor<T> {
   T visitSelectStatement(SelectStatement e);
+  T visitCompoundSelectStatement(CompoundSelectStatement e);
+  T visitCompoundSelectPart(CompoundSelectPart e);
   T visitResultColumn(ResultColumn e);
   T visitInsertStatement(InsertStatement e);
   T visitDeleteStatement(DeleteStatement e);
@@ -274,6 +276,12 @@ class RecursiveVisitor<T> extends AstVisitor<T> {
   @override
   T visitSelectStatement(SelectStatement e) => visitChildren(e);
 
+  @override
+  T visitCompoundSelectStatement(CompoundSelectStatement e) => visitChildren(e);
+
+  @override
+  T visitCompoundSelectPart(CompoundSelectPart e) => visitChildren(e);
+
   @override
   T visitInsertStatement(InsertStatement e) => visitChildren(e);
 

sqlparser/lib/src/ast/statements/select.dart

@@ -1,8 +1,14 @@
 part of '../ast.dart';
 
-class SelectStatement extends Statement
+abstract class BaseSelectStatement extends Statement
-    with CrudStatement, ResultSet
+    with CrudStatement, ResultSet {
-    implements HasWhereClause {
+  /// The resolved list of columns returned by this select statements. Not
+  /// available from the parse tree, will be set later by the analyzer.
+  @override
+  List<Column> resolvedColumns;
+}
+
+class SelectStatement extends BaseSelectStatement implements HasWhereClause {
   final bool distinct;
   final List<ResultColumn> columns;
   final List<Queryable> from;
@@ -15,11 +21,6 @@ class SelectStatement extends Statement
   final OrderByBase orderBy;
   final LimitBase limit;
 
-  /// The resolved list of columns returned by this select statements. Not
-  /// available from the parse tree, will be set later by the analyzer.
-  @override
-  List<Column> resolvedColumns;
-
   SelectStatement(
       {this.distinct = false,
       this.columns,
@@ -54,6 +55,36 @@ class SelectStatement extends Statement
   }
 }
 
+class CompoundSelectStatement extends BaseSelectStatement {
+  final SelectStatement base;
+  final List<CompoundSelectPart> additional;
+
+  // the grammar under https://www.sqlite.org/syntax/compound-select-stmt.html
+  // defines an order by and limit clause on this node, but we parse them as
+  // part of the last compound select statement in [additional]
+
+  CompoundSelectStatement({
+    @required this.base,
+    this.additional = const [],
+  });
+
+  @override
+  Iterable<AstNode> get childNodes {
+    return [base, ...additional];
+  }
+
+  @override
+  T accept<T>(AstVisitor<T> visitor) {
+    return visitor.visitCompoundSelectStatement(this);
+  }
+
+  @override
+  bool contentEquals(CompoundSelectStatement other) {
+    // this class doesn't contain anything but child nodes
+    return true;
+  }
+}
+
 abstract class ResultColumn extends AstNode {
   @override
   T accept<T>(AstVisitor<T> visitor) => visitor.visitResultColumn(this);
@@ -110,3 +141,32 @@ class GroupBy extends AstNode {
     return true; // Defined via child nodes
   }
 }
+
+enum CompoundSelectMode {
+  union,
+  unionAll,
+  intersect,
+  except,
+}
+
+class CompoundSelectPart extends AstNode {
+  final CompoundSelectMode mode;
+  final SelectStatement select;
+
+  /// The first token of this statement, so either union, intersect or except.
+  Token firstModeToken;
+
+  /// The "ALL" token, if this is a "UNION ALL" part
+  Token allToken;
+
+  CompoundSelectPart({@required this.mode, @required this.select});
+
+  @override
+  Iterable<AstNode> get childNodes => [select];
+
+  @override
+  T accept<T>(AstVisitor<T> visitor) => visitor.visitCompoundSelectPart(this);
+
+  @override
+  bool contentEquals(CompoundSelectPart other) => mode == other.mode;
+}

sqlparser/lib/src/reader/parser/crud.dart

@@ -2,7 +2,35 @@ part of 'parser.dart';
 
 mixin CrudParser on ParserBase {
   @override
-  SelectStatement select() {
+  BaseSelectStatement select({bool noCompound}) {
+    if (noCompound == true) {
+      return _selectNoCompound();
+    } else {
+      final first = _selectNoCompound();
+      final parts = <CompoundSelectPart>[];
+
+      while (true) {
+        final part = _compoundSelectPart();
+        if (part != null) {
+          parts.add(part);
+        } else {
+          break;
+        }
+      }
+
+      if (parts.isEmpty) {
+        // no compound parts, just return the simple select statement
+        return first;
+      } else {
+        return CompoundSelectStatement(
+          base: first,
+          additional: parts,
+        )..setSpan(first.first, _previous);
+      }
+    }
+  }
+
+  SelectStatement _selectNoCompound() {
     if (!_match(const [TokenType.select])) return null;
     final selectToken = _previous;
 
@@ -38,6 +66,35 @@ mixin CrudParser on ParserBase {
     )..setSpan(selectToken, _previous);
   }
 
+  CompoundSelectPart _compoundSelectPart() {
+    if (_match(
+        const [TokenType.union, TokenType.intersect, TokenType.except])) {
+      final firstModeToken = _previous;
+      var mode = const {
+        TokenType.union: CompoundSelectMode.union,
+        TokenType.intersect: CompoundSelectMode.intersect,
+        TokenType.except: CompoundSelectMode.except,
+      }[firstModeToken.type];
+      Token allToken;
+
+      if (firstModeToken.type == TokenType.union && _matchOne(TokenType.all)) {
+        allToken = _previous;
+        mode = CompoundSelectMode.unionAll;
+      }
+
+      final select = _selectNoCompound();
+
+      return CompoundSelectPart(
+        mode: mode,
+        select: select,
+      )
+        ..firstModeToken = firstModeToken
+        ..allToken = allToken
+        ..setSpan(firstModeToken, _previous);
+    }
+    return null;
+  }
+
   /// Parses a [ResultColumn] or throws if none is found.
   /// https://www.sqlite.org/syntax/result-column.html
   ResultColumn _resultColumn() {
@@ -114,7 +171,7 @@ mixin CrudParser on ParserBase {
     if (tableRef != null) {
       return tableRef;
     } else if (_matchOne(TokenType.leftParen)) {
-      final innerStmt = select();
+      final innerStmt = _selectNoCompound();
       _consume(TokenType.rightParen,
           'Expected a right bracket to terminate the inner select');
 
@@ -475,7 +532,7 @@ mixin CrudParser on ParserBase {
       _consume(TokenType.$values, 'Expected DEFAULT VALUES');
       return const DefaultValues();
     } else {
-      return SelectInsertSource(select());
+      return SelectInsertSource(_selectNoCompound());
     }
   }
 

sqlparser/lib/src/reader/parser/expressions.dart

@@ -194,7 +194,7 @@ mixin ExpressionParser on ParserBase {
       final existsToken = _previous;
       _consume(
           TokenType.leftParen, 'Expected opening parenthesis after EXISTS');
-      final selectStmt = select();
+      final selectStmt = select(noCompound: true) as SelectStatement;
       _consume(TokenType.rightParen,
           'Expected closing paranthesis to finish EXISTS expression');
       return ExistsExpression(select: selectStmt)
@@ -264,7 +264,7 @@ mixin ExpressionParser on ParserBase {
     if (_matchOne(TokenType.leftParen)) {
       final left = _previous;
       if (_peek.type == TokenType.select) {
-        final stmt = select();
+        final stmt = select(noCompound: true) as SelectStatement;
         _consume(TokenType.rightParen, 'Expected a closing bracket');
         return SubQuery(select: stmt)..setSpan(left, _previous);
       } else {
@@ -379,7 +379,7 @@ mixin ExpressionParser on ParserBase {
         _consume(TokenType.leftParen, 'Expected opening parenthesis for tuple');
     final expressions = <Expression>[];
 
-    final subQuery = select();
+    final subQuery = select(noCompound: true) as SelectStatement;
     if (subQuery == null) {
       // no sub query found. read expressions that form the tuple.
       // tuples can be empty `()`, so only start parsing values when it's not

sqlparser/lib/src/reader/parser/parser.dart

@@ -146,12 +146,13 @@ abstract class ParserBase {
   /// (in brackets) will be accepted as well.
   Expression _consumeTuple({bool orSubQuery = false});
 
-  /// Parses a [SelectStatement], or returns null if there is no select token
+  /// Parses a [BaseSelectStatement], which is either a [SelectStatement] or a
-  /// after the current position.
+  /// [CompoundSelectStatement]. If [noCompound] is set to true, the parser will
+  /// only attempt to parse a [SelectStatement].
   ///
   /// See also:
   /// https://www.sqlite.org/lang_select.html
-  SelectStatement select();
+  BaseSelectStatement select({bool noCompound});
 
   Literal _literalOrNull();
   OrderingMode _orderingModeOrNull();

sqlparser/lib/src/reader/tokenizer/token.dart

@@ -114,6 +114,10 @@ enum TokenType {
   ignore,
   set,
 
+  union,
+  intersect,
+  except,
+
   create,
   table,
   $if,
@@ -237,6 +241,9 @@ const Map<String, TokenType> keywords = {
   'TIES': TokenType.ties,
   'WINDOW': TokenType.window,
   'VALUES': TokenType.$values,
+  'UNION': TokenType.union,
+  'INTERSECT': TokenType.intersect,
+  'EXCEPT': TokenType.except,
 };
 
 /// Maps [TokenType]s which are keywords to their lexeme.

sqlparser/test/parser/select/compound_test.dart

@@ -0,0 +1,40 @@
+import 'package:sqlparser/sqlparser.dart';
+import 'package:test/test.dart';
+
+import '../utils.dart';
+
+void main() {
+  test('parses compound select statements', () {
+    testStatement(
+      'SELECT * FROM tbl UNION ALL SELECT 1 EXCEPT SELECT 2',
+      CompoundSelectStatement(
+        base: SelectStatement(
+          columns: [StarResultColumn(null)],
+          from: [TableReference('tbl', null)],
+        ),
+        additional: [
+          CompoundSelectPart(
+            mode: CompoundSelectMode.unionAll,
+            select: SelectStatement(
+              columns: [
+                ExpressionResultColumn(
+                  expression: NumericLiteral(1, token(TokenType.numberLiteral)),
+                ),
+              ],
+            ),
+          ),
+          CompoundSelectPart(
+            mode: CompoundSelectMode.except,
+            select: SelectStatement(
+              columns: [
+                ExpressionResultColumn(
+                  expression: NumericLiteral(2, token(TokenType.numberLiteral)),
+                ),
+              ],
+            ),
+          ),
+        ],
+      ),
+    );
+  });
+}