diff --git a/crates/polyglot-sql/src/optimizer/qualify_tables.rs b/crates/polyglot-sql/src/optimizer/qualify_tables.rs index 99f8f00a..6997a430 100644 --- a/crates/polyglot-sql/src/optimizer/qualify_tables.rs +++ b/crates/polyglot-sql/src/optimizer/qualify_tables.rs @@ -974,7 +974,7 @@ mod tests { fn test_can_disable_set_operation_wrapper_normalization() { let options = QualifyTablesOptions::new().with_normalize_set_operation_subqueries(false); let expr = parse( - "SELECT * FROM (SELECT * FROM tab_1) UNION ALL SELECT * FROM (SELECT * FROM tab_1)", + "SELECT * FROM (SELECT * FROM (SELECT * FROM tab_1) UNION ALL SELECT * FROM (SELECT * FROM tab_1))", ); let qualified = qualify_tables(expr, &options); let sql = gen(&qualified); diff --git a/crates/polyglot-sql/src/parser.rs b/crates/polyglot-sql/src/parser.rs index 639a21fb..7535e7a7 100644 --- a/crates/polyglot-sql/src/parser.rs +++ b/crates/polyglot-sql/src/parser.rs @@ -4681,6 +4681,11 @@ impl Parser { })) } } else if self.match_token(TokenType::LParen) { + // A parenthesized query continues into a set operation only where it opens an + // enclosing parenthesis, as in `FROM ((SELECT 1) UNION (SELECT 2))`. In + // `FROM (SELECT ...) UNION ALL SELECT ...` the set operation combines the outer query. + let set_operation_may_follow = + self.current < 2 || self.tokens[self.current - 2].token_type == TokenType::LParen; // Subquery or parenthesized set operation or (VALUES ...) if self.check(TokenType::Values) { // (VALUES (...), (...)) AS t(c1, c2) or (VALUES (0) foo(bar)) @@ -4735,9 +4740,10 @@ impl Parser { // Check for set operations after parenthesized query // If there's a set operation, wrap query in Subquery first to preserve parens // e.g., (SELECT 1) UNION (SELECT 2) - the left operand needs Subquery wrapper - let result = if self.check(TokenType::Union) - || self.check(TokenType::Intersect) - || self.check(TokenType::Except) + let result = if set_operation_may_follow + && (self.check(TokenType::Union) + || self.check(TokenType::Intersect) + || self.check(TokenType::Except)) { let left = Expression::Subquery(Box::new(Subquery { this: query, diff --git a/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs b/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs index fffb9003..149dec7f 100644 --- a/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs +++ b/crates/polyglot-sql/tests/set_operation_subquery_operand_regression.rs @@ -98,6 +98,34 @@ fn set_operation_after_subquery_operand_applies_to_enclosing_query() { } } +#[test] +fn set_operation_after_from_subquery_applies_to_enclosing_query() { + let cases = [ + "SELECT * FROM (SELECT o.id FROM orders AS o) UNION ALL SELECT 1", + "SELECT * FROM (SELECT o.id FROM orders AS o) UNION SELECT 1", + "SELECT * FROM (SELECT o.id FROM orders AS o) EXCEPT SELECT 1", + "SELECT * FROM (SELECT o.id FROM orders AS o) INTERSECT SELECT 1", + "SELECT * FROM (SELECT o.id FROM orders AS o) AS t UNION ALL SELECT 1", + "SELECT * FROM (WITH c AS (SELECT 1 AS id) SELECT c.id FROM c) UNION ALL SELECT 1", + "SELECT * FROM (SELECT o.id FROM orders AS o) UNION ALL SELECT * FROM (SELECT 1 AS id)", + ]; + for dialect in DIALECTS { + for sql in cases { + let parsed = parse_stable(sql, dialect); + let (left, right) = set_operation_operands(&parsed); + assert!( + matches!(left, Expression::Select(_)) && matches!(right, Expression::Select(_)), + "{dialect:?}: {sql}: {parsed:?}" + ); + assert_eq!( + set_operation_count(&parsed), + 1, + "{dialect:?}: set operation absorbed by a FROM subquery: {sql}" + ); + } + } +} + #[test] fn set_operation_inside_enclosing_parentheses_stays_nested() { let cases = [