Merge pull request #98 from supabase/feat/combining-queries

feat: add support for combining queries

Author: psteinroe

Diff for: crates/codegen/src/get_node_properties.rs

@@ -181,6 +181,17 @@ fn custom_handlers(node: &Node) -> TokenStream {
                 tokens.push(TokenProperty::from(Token::GroupP));
                 tokens.push(TokenProperty::from(Token::By));
             }
+            match n.op() {
+                protobuf::SetOperation::Undefined => {},
+                protobuf::SetOperation::SetopNone => {},
+                protobuf::SetOperation::SetopUnion => tokens.push(TokenProperty::from(Token::Union)),
+                protobuf::SetOperation::SetopIntersect => tokens.push(TokenProperty::from(Token::Intersect)),
+                protobuf::SetOperation::SetopExcept => tokens.push(TokenProperty::from(Token::Except)),
+                _ => panic!("Unknown SelectStmt op {:#?}", n.op()),
+            }
+            if n.all {
+                tokens.push(TokenProperty::from(Token::All));
+            }
         },
         "BoolExpr" => quote! {
             match n.boolop() {
@@ -391,6 +402,12 @@ fn custom_handlers(node: &Node) -> TokenStream {
             tokens.push(TokenProperty::from(Token::View));
             if n.query.is_some() {
                 tokens.push(TokenProperty::from(Token::As));
+                // check if SelectStmt with WithClause with recursive set to true
+                if let Some(NodeEnum::SelectStmt(select_stmt)) = n.query.as_ref().and_then(|query| query.node.as_ref()) {
+                    if select_stmt.with_clause.is_some() && select_stmt.with_clause.as_ref().unwrap().recursive {
+                        tokens.push(TokenProperty::from(Token::Recursive));
+                    }
+                }
             }
             if n.replace {
                 tokens.push(TokenProperty::from(Token::Or));

Diff for: crates/parser/src/parse/statement.rs

@@ -74,6 +74,11 @@ fn collect_statement_token_range(parser: &mut Parser, kind: SyntaxKind) -> Range
     let mut ignore_next_non_whitespace = false;
     while !parser.at(SyntaxKind::Ascii59) && !parser.eof() {
         match parser.nth(0, false).kind {
+            SyntaxKind::All => {
+                // ALL is never a statement start, but needs to be skipped when combining queries
+                // (e.g. UNION ALL)
+                parser.advance();
+            }
             SyntaxKind::BeginP => {
                 // BEGIN, consume until END
                 is_sub_trx += 1;
@@ -92,7 +97,7 @@ fn collect_statement_token_range(parser: &mut Parser, kind: SyntaxKind) -> Range
                 is_sub_stmt -= 1;
                 parser.advance();
             }
-            SyntaxKind::As => {
+            SyntaxKind::As | SyntaxKind::Union | SyntaxKind::Intersect | SyntaxKind::Except => {
                 // ignore the next non-whitespace token
                 ignore_next_non_whitespace = true;
                 parser.advance();

Diff for: crates/parser/tests/data/statements/valid/0044.sql

@@ -2,4 +2,3 @@ CREATE VIEW comedies AS SELECT * FROM films WHERE kind = 'Comedy';
 CREATE VIEW universal_comedies AS SELECT * FROM comedies WHERE classification = 'U' WITH LOCAL CHECK OPTION;
 CREATE VIEW pg_comedies AS SELECT * FROM comedies WHERE classification = 'PG' WITH CASCADED CHECK OPTION;
 CREATE VIEW comedies AS SELECT f.*, country_code_to_name(f.country_code) AS country, (SELECT avg(r.rating) FROM user_ratings r WHERE r.film_id = f.id) AS avg_rating FROM films f WHERE f.kind = 'Comedy';
-CREATE RECURSIVE VIEW public.nums_1_100 (n) AS VALUES (1) UNION ALL SELECT n+1 FROM nums_1_100 WHERE n < 100;

Diff for: crates/parser/tests/data/statements/valid/0055.sql

@@ -0,0 +1,2 @@
+select 1 union all select 2;
+select 1 union select 2;