Implement UNION / UNION ALL

beanquery/parser/ast.py

@@ -81,11 +81,22 @@ def node(name, fields):
 #   from_clause: An instance of 'From', or None if absent.
 #   where_clause: A root expression node, or None if absent.
 #   group_by: An instance of 'GroupBy', or None if absent.
-#   order_by: An instance of 'OrderBy', or None if absent.
-#   pivot_by: An instance of 'PivotBy', or None if absent.
-#   limit: An integer, or None is absent.
 #   distinct: A boolean value (True), or None if absent.
-Select = node('Select', 'targets from_clause where_clause group_by order_by pivot_by limit distinct')
+Select = node('Select', 'targets from_clause where_clause group_by distinct')
+
+# The top-level query node wrapping one or more SELECT bodies.
+#
+# A single SELECT is the degenerate case (len(queries) == 1).
+# In the future, UNION chain support will be added where len(queries) > 1.
+#
+# Attributes:
+#   queries:   List of Select nodes.
+#   set_operators: List of set-operator names between adjacent queries, e.g.
+#              'union' or 'union_all'.  len == len(queries) - 1.
+#   order_by:  Optional list of OrderBy applied to the combined result.
+#   limit:     Optional integer limit applied to the combined result.
+#   pivot_by:  Optional PivotBy applied to the combined result.
+Query = node('Query', 'queries set_operators order_by limit pivot_by')
 
 # A select query that produces final balances for accounts.
 # This is equivalent to

beanquery/parser/bql.ebnf

@@ -15,26 +15,42 @@ bql
 
 statement
     =
-    | select
+    | query
     | balances
     | journal
     | print
     | create_table
     | insert
     ;
 
+(* Wrapper for queries, which can be either a plain SELECT or a chain of
+   SELECTs combined by set operators (UNION, UNION ALL).  ORDER BY, LIMIT,
+   and PIVOT BY after the last operand apply to the combined result set.
+   The list `operators` has one entry per inter-operand connector, always one
+   less than the number of queries. *)
+query::Query
+    = queries+:( select | subquery )
+      { ( 'UNION' 'ALL' set_operators+:`'union_all'`
+        | 'UNION'       set_operators+:`'union'`
+        ) queries+:( select | subquery)
+      }*
+      ['ORDER' 'BY' order_by:','.{order}+]
+      ['LIMIT' limit:integer]
+      ['PIVOT' 'BY' pivot_by:pivotby]
+    ;
+
+(* SELECT body without ORDER BY / LIMIT / PIVOT BY so the enclosing query rule can claim
+   those tokens for the combined result set. *)
 select::Select
     = 'SELECT' ['DISTINCT' distinct:`True`] targets:(','.{ target }+ | asterisk)
-      ['FROM' from_clause:(_table | subselect | from)]
+      ['FROM' from_clause:(_table | subquery | from)]
       ['WHERE' where_clause:expression]
       ['GROUP' 'BY' group_by:groupby]
-      ['ORDER' 'BY' order_by:','.{order}+]
-      ['PIVOT' 'BY' pivot_by:pivotby]
-      ['LIMIT' limit:integer]
     ;
 
-subselect
-    = '(' @:select ')'
+(* Parenthesised sub-query; uses query so ORDER BY / LIMIT are allowed inside. *)
+subquery
+    = '(' @:query ')'
     ;
 
 from::From
@@ -132,12 +148,20 @@ comparison
     | sum
     ;
 
+(* This operator is special in that it has parentheses. Avoid double parentheses
+with subquerys ALL( (SELECT ...) ) by &(...) look-ahead *)
 any::Any
-    = left:sum op:op 'any' '(' right:expression ')'
+    =
+    | left:sum op:op 'any' &('(' 'SELECT') right:subquery
+    | left:sum op:op 'any' '(' right:expression ')'
     ;
 
+(* This operator is special in that it has parentheses. Avoid double parentheses
+with subquerys ALL( (SELECT ...) ) by &(...) look-ahead *)
 all::All
-    = left:sum op:op 'all' '(' right:expression ')'
+    =
+    | left:sum op:op 'all' &('(' 'SELECT') right:subquery
+    | left:sum op:op 'all' '(' right:expression ')'
     ;
 
 op
@@ -282,7 +306,7 @@ subscript::Subscript
 
 atom
     =
-    | select
+    | subquery
     | function
     | constant
     | column
@@ -390,7 +414,7 @@ create_table::CreateTable
       (
       | '(' columns:','.{( identifier identifier )} ')' ['USING' using:string]
       | 'USING' using:string
-      | 'AS' query:select
+      | 'AS' query:query
       )
     ;