add UNION query support by concatenating results in python

django_ormql/engine.py

@@ -16,11 +16,16 @@ def query(
         placeholders=None,
         timezone=datetime.timezone.utc,
         default_limit=None,
+        dry_run=False,
     ):
-        return Query(
+        query = Query(
             query,
             self.tables,
             placeholders,
             timezone,
             default_limit,
-        ).evaluate()
+        )
+        if dry_run:
+            return query.parse()
+        else:
+            return query.evaluate()

django_ormql/query.py

@@ -48,6 +48,7 @@ class Tokenizer(Tokenizer):
             "!=": TokenType.NEQ,
             "||": TokenType.DPIPE,
             "->": TokenType.ARROW,
+            "ALL": TokenType.ALL,
             "AND": TokenType.AND,
             "ASC": TokenType.ASC,
             "AS": TokenType.ALIAS,
@@ -82,6 +83,7 @@ class Tokenizer(Tokenizer):
             "SELECT": TokenType.SELECT,
             "THEN": TokenType.THEN,
             "TRUE": TokenType.TRUE,
+            "UNION": TokenType.UNION,
             "WHEN": TokenType.WHEN,
             "WHERE": TokenType.WHERE,
             # TYPES
@@ -657,6 +659,9 @@ def _where_to_django(self, node, **kwargs):
         return self._expression_to_django(node, **kwargs)
 
     def _select_to_qs(self, root, parent_table_stack):
+        if not isinstance(root, expressions.Select):
+            raise QueryNotSupported("Only SELECT queries are supported")
+
         table = root.args["from_"].this
         if not isinstance(table, expressions.Table):
             raise QueryNotSupported("Unsupported FROM statement")
@@ -838,7 +843,19 @@ def _select_to_qs(self, root, parent_table_stack):
 
         return qs, values_names
 
-    def evaluate(self):
+    def _flatten_unions(self, root):
+        if isinstance(root, expressions.Select):
+            return [root]
+        elif isinstance(root, expressions.Subquery) and isinstance(root.this, expressions.Select):
+            return [root.this]
+        elif isinstance(root, expressions.Union) and not root.args['distinct']:
+            if root.args.get('limit') or root.args.get('order') or root.args.get('offset'):
+                raise QueryError("ORDER, LIMIT and OFFSET modifiers are not supported on UNION queries")
+            return self._flatten_unions(root.left) + self._flatten_unions(root.right)
+        else:
+            raise QueryNotSupported("Only SELECT and SELECT ... UNION ALL queries are supported")
+
+    def parse(self):
         try:
             ast = parse_one(self.sql, dialect=OrmqlDialect)
         except ParseError as e:
@@ -848,27 +865,33 @@ def evaluate(self):
         if settings.DEBUG:
             print(f"Parsed statement: {ast!r}")
 
-        if not isinstance(ast, expressions.Select):
-            raise QueryNotSupported("Only SELECT queries are supported")
-
         try:
-            qs, values_names = self._select_to_qs(ast, [])
+            queries = self._flatten_unions(ast)
+            results = [self._select_to_qs(query, []) for query in queries]
+            if len(set(len(values_names.keys()) for qs, values_names in results)) != 1:
+                raise QueryError("All parts of UNION query must return same number of columns")
         except QueryError:
             raise
         except FieldError as e:
             raise QueryError("Invalid combination of types") from e
         except Exception as e:
             raise QueryError("Query parsing failed") from e
 
-        if isinstance(qs, dict):
-            yield {values_names[k]: v for k, v in qs.items()}
-        else:
-            try:
-                if settings.DEBUG:
-                    print(f"Generated statement: {qs.query!s}")
-                for row in qs:
-                    yield {
-                        values_names[k]: v for k, v in row.items() if k in values_names
-                    }
-            except (FieldError, ValueError) as e:
-                raise QueryError("Invalid combination of types") from e
+        return [qs for qs, values_names in results], results[0][1]
+
+    def evaluate(self):
+        querysets, values_names = self.parse()
+
+        for qs in querysets:
+            if isinstance(qs, dict):
+                yield {values_names[k]: v for k, v in qs.items()}
+            else:
+                try:
+                    if settings.DEBUG:
+                        print(f"Generated statement: {qs.query!s}")
+                    for row in qs:
+                        yield {
+                            values_names[k]: v for k, v in row.items() if k in values_names
+                        }
+                except (FieldError, ValueError) as e:
+                    raise QueryError("Invalid combination of types") from e

doc/syntax.md

@@ -17,7 +17,8 @@
   - Bitwise operators
   - `IS DISTINCT`
   - `LIMIT` and `OFFSET` with complex expressions
-  - `UNION`, `INTERSECT`, `EXCEPT`
+  - `UNION` queries without `ALL`, `UNION` with `ORDER BY`, `LIMIT`, `OFFSET` modifiers 
+  - `INTERSECT`, `EXCEPT`
 
 This list of differences is not complete.
 

tests/test_select_col_from.py

@@ -234,7 +234,7 @@ def test_undeclared_field(engine_t1):
 
 @pytest.mark.django_db
 def test_compound_not_allowed(engine_t1):
-    with pytest.raises(QueryError, match="Invalid expression / Unexpected token"):
+    with pytest.raises(QueryError, match="Only SELECT and SELECT ... UNION ALL queries are supported"):
         list(
             engine_t1.query(
                 """
@@ -274,7 +274,7 @@ def test_compound_not_allowed(engine_t1):
 
 @pytest.mark.django_db
 def test_values_query_not_allowed(engine_t1):
-    with pytest.raises(QueryError, match="Only SELECT queries are supported"):
+    with pytest.raises(QueryError, match="Only SELECT and SELECT ... UNION ALL queries are supported"):
         list(
             engine_t1.query(
                 """

tests/test_subqueries.py

@@ -202,3 +202,17 @@ def test_invalid_outerref(engine_t1):
             """
             )
         )
+
+
+@pytest.mark.django_db
+def test_union_subquery_not_allowed(engine_t1):
+    with pytest.raises(QueryError, match="Only SELECT subqueries are supported"):
+        list(
+            engine_t1.query(
+                """
+            SELECT title
+            FROM products
+            WHERE EXISTS(SELECT 1 FROM orderpositions WHERE product = OUTER(id) AND order.status = "paid" UNION SELECT 1 FROM orderpositions WHERE product = OUTER(id) AND order.status = "canceled")
+            """
+            )
+        )

tests/test_union.py

@@ -0,0 +1,151 @@
+from decimal import Decimal
+
+import pytest
+
+from django_ormql.exceptions import QueryNotSupported, QueryError
+
+
+@pytest.mark.django_db
+def test_union_modifiers_not_allowed(engine_t1):
+    with pytest.raises(QueryError, match="Only SELECT and SELECT ... UNION ALL queries are supported"):
+        list(
+            engine_t1.query(
+                """
+            SELECT title
+            FROM categories
+            UNION
+            SELECT title
+            FROM products
+            """
+            )
+        )
+    with pytest.raises(QueryError, match="ORDER, LIMIT and OFFSET modifiers are not supported on UNION queries"):
+        list(
+            engine_t1.query(
+                """
+            SELECT title
+            FROM categories
+            UNION ALL
+            SELECT title
+            FROM products
+            LIMIT 1
+            """
+            )
+        )
+    with pytest.raises(QueryError, match="ORDER, LIMIT and OFFSET modifiers are not supported on UNION queries"):
+        list(
+            engine_t1.query(
+                """
+            SELECT title
+            FROM categories
+            UNION ALL
+            SELECT title
+            FROM products
+            ORDER BY title
+            """
+            )
+        )
+
+
+@pytest.mark.django_db
+def test_union(engine_t1):
+    res = engine_t1.query(
+        """
+        SELECT 'Category' type, title
+        FROM categories
+        UNION ALL
+        SELECT 'Product' type, title
+        FROM products
+        """
+    )
+    assert list(res) == [
+        {"type": "Category", "title": "Books"},
+        {"type": "Category", "title": "DVDs"},
+        {"type": "Product", "title": "Lord of the rings"},
+        {"type": "Product", "title": "SQL for Dummies"},
+        {"type": "Product", "title": "Lord of the rings DVD"},
+    ]
+    res = engine_t1.query(
+        """
+        (SELECT 'Category' type, title
+        FROM categories ORDER BY title DESC)
+        UNION ALL
+        (SELECT 'Product' type, title
+        FROM products ORDER BY title DESC)
+        """
+    )
+    assert list(res) == [
+        {"type": "Category", "title": "DVDs"},
+        {"type": "Category", "title": "Books"},
+        {"type": "Product", "title": "SQL for Dummies"},
+        {"type": "Product", "title": "Lord of the rings DVD"},
+        {"type": "Product", "title": "Lord of the rings"},
+    ]
+
+
+@pytest.mark.django_db
+def test_union_col_count(engine_t1):
+    with pytest.raises(QueryError, match="All parts of UNION query must return same number of columns"):
+        list(
+            engine_t1.query(
+                """
+            SELECT title
+            FROM categories
+            UNION ALL
+            SELECT category.title, SUM(price)
+            FROM products
+            GROUP BY category.title
+            """
+            )
+        )
+    with pytest.raises(QueryError, match="All parts of UNION query must return same number of columns"):
+        list(
+            engine_t1.query(
+                """
+            SELECT title
+            FROM categories
+            UNION ALL
+            SELECT title, price
+            FROM products
+            """
+            )
+        )
+
+
+@pytest.mark.django_db
+def test_union_grouped(engine_t1):
+    res = engine_t1.query(
+        """
+        SELECT 'Category' type, title, 0 sales, 0 total
+        FROM categories
+        UNION ALL
+        SELECT 'Product' type, product.title, SUM(quantity) sales, SUM(single_price * quantity) total
+        FROM orderpositions
+        GROUP BY product.id
+        """
+    )
+    assert list(res) == [
+        {"type": "Category", "title": "Books", "sales": 0, "total": 0},
+        {"type": "Category", "title": "DVDs", "sales": 0, "total": 0},
+        {"type": "Product", "title": "Lord of the rings", "sales": 3, "total": Decimal('32.10')},
+        {"type": "Product", "title": "SQL for Dummies", "sales": 1, "total": Decimal('21.40')},
+        {"type": "Product", "title": "Lord of the rings DVD", "sales": 4, "total": Decimal('76.00')},
+    ]
+    res = engine_t1.query(
+        """
+        SELECT 'Category' type, product.category.title title, SUM(quantity) sales, SUM(single_price * quantity) total
+        FROM orderpositions
+        GROUP BY product.category.id
+        UNION ALL
+        SELECT 'Product' type, product.title, SUM(quantity) sales, SUM(single_price * quantity) total
+        FROM orderpositions
+        GROUP BY product.id
+        """
+    )
+    assert list(res) == [
+        {"type": "Category", "title": "Books", "sales": 4, "total": Decimal('53.50')},
+        {"type": "Category", "title": "DVDs", "sales": 4, "total": Decimal('76.00')},
+        {"type": "Product", "title": "Lord of the rings", "sales": 3, "total": Decimal('32.10')},
+        {"type": "Product", "title": "SQL for Dummies", "sales": 1, "total": Decimal('21.40')},
+        {"type": "Product", "title": "Lord of the rings DVD", "sales": 4, "total": Decimal('76.00')},
+    ]