diff --git a/django_ormql/engine.py b/django_ormql/engine.py index d4eb9cd..d5807c6 100644 --- a/django_ormql/engine.py +++ b/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() diff --git a/django_ormql/query.py b/django_ormql/query.py index 3bcb9db..72bc504 100644 --- a/django_ormql/query.py +++ b/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,11 +865,11 @@ 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: @@ -860,15 +877,21 @@ def evaluate(self): 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 diff --git a/doc/syntax.md b/doc/syntax.md index d11a5e3..eb6e7fd 100644 --- a/doc/syntax.md +++ b/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. diff --git a/tests/test_select_col_from.py b/tests/test_select_col_from.py index 9e1839d..e41981c 100644 --- a/tests/test_select_col_from.py +++ b/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( """ diff --git a/tests/test_subqueries.py b/tests/test_subqueries.py index d35446f..12d483f 100644 --- a/tests/test_subqueries.py +++ b/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") + """ + ) + ) diff --git a/tests/test_union.py b/tests/test_union.py new file mode 100644 index 0000000..7b2e89c --- /dev/null +++ b/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')}, + ] \ No newline at end of file