From f1b2356184bfc08c9ec060e7fadb15ce8c138160 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Fri, 29 Dec 2023 08:57:01 +0100 Subject: [PATCH] feat: add support for combining queries ref: https://www.postgresql.org/docs/current/queries-union.html also fixed some minor bugs with `ViewStmt` --- crates/codegen/src/get_node_properties.rs | 17 + crates/parser/src/parse/statement.rs | 7 +- .../tests/data/statements/valid/0044.sql | 1 - .../tests/data/statements/valid/0055.sql | 2 + .../snapshots/statements/valid/0044@5.snap | 619 ++++++++++++------ .../snapshots/statements/valid/0055@1.snap | 165 +++++ .../snapshots/statements/valid/0055@2.snap | 163 +++++ 7 files changed, 776 insertions(+), 198 deletions(-) create mode 100644 crates/parser/tests/data/statements/valid/0055.sql create mode 100644 crates/parser/tests/snapshots/statements/valid/0055@1.snap create mode 100644 crates/parser/tests/snapshots/statements/valid/0055@2.snap diff --git a/crates/codegen/src/get_node_properties.rs b/crates/codegen/src/get_node_properties.rs index b42caf3b..0adf77fa 100644 --- a/crates/codegen/src/get_node_properties.rs +++ b/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 --git a/crates/parser/src/parse/statement.rs b/crates/parser/src/parse/statement.rs index 150c45e3..e1a7c5ff 100644 --- a/crates/parser/src/parse/statement.rs +++ b/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 --git a/crates/parser/tests/data/statements/valid/0044.sql b/crates/parser/tests/data/statements/valid/0044.sql index 26334f8c..68998681 100644 --- a/crates/parser/tests/data/statements/valid/0044.sql +++ b/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 --git a/crates/parser/tests/data/statements/valid/0055.sql b/crates/parser/tests/data/statements/valid/0055.sql new file mode 100644 index 00000000..6f05a39b --- /dev/null +++ b/crates/parser/tests/data/statements/valid/0055.sql @@ -0,0 +1,2 @@ +select 1 union all select 2; +select 1 union select 2; diff --git a/crates/parser/tests/snapshots/statements/valid/0044@5.snap b/crates/parser/tests/snapshots/statements/valid/0044@5.snap index e5bc0e53..1b163db8 100644 --- a/crates/parser/tests/snapshots/statements/valid/0044@5.snap +++ b/crates/parser/tests/snapshots/statements/valid/0044@5.snap @@ -3,101 +3,112 @@ source: crates/parser/tests/statement_parser_test.rs description: CREATE RECURSIVE VIEW public.nums_1_100 (n) AS VALUES (1) UNION ALL SELECT n+1 FROM nums_1_100 WHERE n < 100; --- Parse { - cst: SourceFile@0..108 - Create@0..6 "CREATE" - Whitespace@6..7 " " - Recursive@7..16 "RECURSIVE" - Whitespace@16..17 " " - View@17..21 "VIEW" - Whitespace@21..22 " " - Ident@22..28 "public" - Ascii46@28..29 "." - Ident@29..39 "nums_1_100" - Whitespace@39..40 " " - Ascii40@40..41 "(" - Ident@41..42 "n" - Ascii41@42..43 ")" - Whitespace@43..44 " " - As@44..46 "AS" - Whitespace@46..47 " " - Values@47..53 "VALUES" - Whitespace@53..54 " " - Ascii40@54..55 "(" - Iconst@55..56 "1" - Ascii41@56..57 ")" - Whitespace@57..58 " " - Union@58..63 "UNION" - Whitespace@63..64 " " - All@64..67 "ALL" - SelectStmt@67..108 - Select@67..73 "SELECT" - Whitespace@73..74 " " - ResTarget@74..77 - AExpr@74..77 - ColumnRef@74..75 - Ident@74..75 "n" - Ascii43@75..76 "+" - AConst@76..77 - Iconst@76..77 "1" - Whitespace@77..78 " " - From@78..82 "FROM" - Whitespace@82..83 " " - RangeVar@83..93 - Ident@83..93 "nums_1_100" - Whitespace@93..94 " " - Where@94..99 "WHERE" - Whitespace@99..100 " " - AExpr@100..107 - ColumnRef@100..101 - Ident@100..101 "n" - Whitespace@101..102 " " - Ascii60@102..103 "<" - Whitespace@103..104 " " - AConst@104..107 - Iconst@104..107 "100" - Ascii59@107..108 ";" + cst: SourceFile@0..109 + ViewStmt@0..109 + Create@0..6 "CREATE" + Whitespace@6..7 " " + Recursive@7..16 "RECURSIVE" + Whitespace@16..17 " " + View@17..21 "VIEW" + Whitespace@21..22 " " + RangeVar@22..39 + Ident@22..28 "public" + Ascii46@28..29 "." + Ident@29..39 "nums_1_100" + Whitespace@39..40 " " + Ascii40@40..41 "(" + Ident@41..42 "n" + Ascii41@42..43 ")" + Whitespace@43..44 " " + As@44..46 "AS" + Whitespace@46..47 " " + SelectStmt@47..109 + WithClause@47..108 + CommonTableExpr@47..108 + SelectStmt@47..108 + SelectStmt@47..56 + Values@47..53 "VALUES" + Whitespace@53..54 " " + Ascii40@54..55 "(" + List@55..56 + AConst@55..56 + Iconst@55..56 "1" + Ascii41@56..57 ")" + Whitespace@57..58 " " + Union@58..63 "UNION" + Whitespace@63..64 " " + All@64..67 "ALL" + Whitespace@67..68 " " + Select@68..74 "SELECT" + Whitespace@74..75 " " + SelectStmt@75..108 + ResTarget@75..78 + AExpr@75..78 + ColumnRef@75..76 + Ident@75..76 "n" + Ascii43@76..77 "+" + AConst@77..78 + Iconst@77..78 "1" + Whitespace@78..79 " " + From@79..83 "FROM" + Whitespace@83..84 " " + RangeVar@84..94 + Ident@84..94 "nums_1_100" + Whitespace@94..95 " " + Where@95..100 "WHERE" + Whitespace@100..101 " " + AExpr@101..108 + ColumnRef@101..102 + Ident@101..102 "n" + Whitespace@102..103 " " + Ascii60@103..104 "<" + Whitespace@104..105 " " + AConst@105..108 + Iconst@105..108 "100" + Ascii59@108..109 ";" , - errors: [ - SyntaxError( - "Expected Ascii59, found Whitespace", - 67..67, - ), - SyntaxError( - "Invalid statement: syntax error at end of input", - 0..25, - ), - ], + errors: [], stmts: [ RawStmt { - stmt: SelectStmt( - SelectStmt { - distinct_clause: [], - into_clause: None, - target_list: [ + stmt: ViewStmt( + ViewStmt { + view: Some( + RangeVar { + catalogname: "", + schemaname: "public", + relname: "nums_1_100", + inh: true, + relpersistence: "p", + alias: None, + location: 22, + }, + ), + aliases: [ + Node { + node: Some( + String( + String { + sval: "n", + }, + ), + ), + }, + ], + query: Some( Node { node: Some( - ResTarget( - ResTarget { - name: "", - indirection: [], - val: Some( + SelectStmt( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [ Node { node: Some( - AExpr( - AExpr { - kind: AexprOp, - name: [ - Node { - node: Some( - String( - String { - sval: "+", - }, - ), - ), - }, - ], - lexpr: Some( + ResTarget( + ResTarget { + name: "", + indirection: [], + val: Some( Node { node: Some( ColumnRef( @@ -113,142 +124,358 @@ Parse { ), }, ], - location: 7, - }, - ), - ), - }, - ), - rexpr: Some( - Node { - node: Some( - AConst( - AConst { - isnull: false, - location: 9, - val: Some( - Ival( - Integer { - ival: 1, - }, - ), - ), + location: -1, }, ), ), }, ), - location: 8, + location: -1, }, ), ), }, - ), - location: 7, - }, - ), - ), - }, - ], - from_clause: [ - Node { - node: Some( - RangeVar( - RangeVar { - catalogname: "", - schemaname: "", - relname: "nums_1_100", - inh: true, - relpersistence: "p", - alias: None, - location: 16, - }, - ), - ), - }, - ], - where_clause: Some( - Node { - node: Some( - AExpr( - AExpr { - kind: AexprOp, - name: [ + ], + from_clause: [ Node { node: Some( - String( - String { - sval: "<", + RangeVar( + RangeVar { + catalogname: "", + schemaname: "", + relname: "nums_1_100", + inh: true, + relpersistence: "p", + alias: None, + location: -1, }, ), ), }, ], - lexpr: Some( - Node { - node: Some( - ColumnRef( - ColumnRef { - fields: [ - Node { - node: Some( - String( - String { - sval: "n", - }, - ), + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: Some( + WithClause { + ctes: [ + Node { + node: Some( + CommonTableExpr( + CommonTableExpr { + ctename: "nums_1_100", + aliascolnames: [ + Node { + node: Some( + String( + String { + sval: "n", + }, + ), + ), + }, + ], + ctematerialized: Default, + ctequery: Some( + Node { + node: Some( + SelectStmt( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopUnion, + all: true, + larg: Some( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [ + Node { + node: Some( + List( + List { + items: [ + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 55, + val: Some( + Ival( + Integer { + ival: 1, + }, + ), + ), + }, + ), + ), + }, + ], + }, + ), + ), + }, + ], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopNone, + all: false, + larg: None, + rarg: None, + }, + ), + rarg: Some( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [ + Node { + node: Some( + ResTarget( + ResTarget { + name: "", + indirection: [], + val: Some( + Node { + node: Some( + AExpr( + AExpr { + kind: AexprOp, + name: [ + Node { + node: Some( + String( + String { + sval: "+", + }, + ), + ), + }, + ], + lexpr: Some( + Node { + node: Some( + ColumnRef( + ColumnRef { + fields: [ + Node { + node: Some( + String( + String { + sval: "n", + }, + ), + ), + }, + ], + location: 75, + }, + ), + ), + }, + ), + rexpr: Some( + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 77, + val: Some( + Ival( + Integer { + ival: 1, + }, + ), + ), + }, + ), + ), + }, + ), + location: 76, + }, + ), + ), + }, + ), + location: 75, + }, + ), + ), + }, + ], + from_clause: [ + Node { + node: Some( + RangeVar( + RangeVar { + catalogname: "", + schemaname: "", + relname: "nums_1_100", + inh: true, + relpersistence: "p", + alias: None, + location: 84, + }, + ), + ), + }, + ], + where_clause: Some( + Node { + node: Some( + AExpr( + AExpr { + kind: AexprOp, + name: [ + Node { + node: Some( + String( + String { + sval: "<", + }, + ), + ), + }, + ], + lexpr: Some( + Node { + node: Some( + ColumnRef( + ColumnRef { + fields: [ + Node { + node: Some( + String( + String { + sval: "n", + }, + ), + ), + }, + ], + location: 101, + }, + ), + ), + }, + ), + rexpr: Some( + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 105, + val: Some( + Ival( + Integer { + ival: 100, + }, + ), + ), + }, + ), + ), + }, + ), + location: 103, + }, + ), + ), + }, + ), + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopNone, + all: false, + larg: None, + rarg: None, + }, + ), + }, + ), + ), + }, ), + search_clause: None, + cycle_clause: None, + location: -1, + cterecursive: false, + cterefcount: 0, + ctecolnames: [], + ctecoltypes: [], + ctecoltypmods: [], + ctecolcollations: [], }, - ], - location: 33, - }, - ), - ), - }, - ), - rexpr: Some( - Node { - node: Some( - AConst( - AConst { - isnull: false, - location: 37, - val: Some( - Ival( - Integer { - ival: 100, - }, - ), ), - }, - ), - ), + ), + }, + ], + recursive: true, + location: -1, }, ), - location: 35, + op: SetopNone, + all: false, + larg: None, + rarg: None, }, ), ), }, ), - group_clause: [], - group_distinct: false, - having_clause: None, - window_clause: [], - values_lists: [], - sort_clause: [], - limit_offset: None, - limit_count: None, - limit_option: Default, - locking_clause: [], - with_clause: None, - op: SetopNone, - all: false, - larg: None, - rarg: None, + replace: false, + options: [], + with_check_option: NoCheckOption, }, ), - range: 67..108, + range: 0..108, }, ], } diff --git a/crates/parser/tests/snapshots/statements/valid/0055@1.snap b/crates/parser/tests/snapshots/statements/valid/0055@1.snap new file mode 100644 index 00000000..5fd77413 --- /dev/null +++ b/crates/parser/tests/snapshots/statements/valid/0055@1.snap @@ -0,0 +1,165 @@ +--- +source: crates/parser/tests/statement_parser_test.rs +description: select 1 union all select 2; +--- +Parse { + cst: SourceFile@0..28 + SelectStmt@0..28 + Select@0..6 "select" + Whitespace@6..7 " " + SelectStmt@7..8 + ResTarget@7..8 + AConst@7..8 + Iconst@7..8 "1" + Whitespace@8..9 " " + Union@9..14 "union" + Whitespace@14..15 " " + All@15..18 "all" + Whitespace@18..19 " " + SelectStmt@19..27 + Select@19..25 "select" + Whitespace@25..26 " " + ResTarget@26..27 + AConst@26..27 + Iconst@26..27 "2" + Ascii59@27..28 ";" + , + errors: [], + stmts: [ + RawStmt { + stmt: SelectStmt( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopUnion, + all: true, + larg: Some( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [ + Node { + node: Some( + ResTarget( + ResTarget { + name: "", + indirection: [], + val: Some( + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 7, + val: Some( + Ival( + Integer { + ival: 1, + }, + ), + ), + }, + ), + ), + }, + ), + location: 7, + }, + ), + ), + }, + ], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopNone, + all: false, + larg: None, + rarg: None, + }, + ), + rarg: Some( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [ + Node { + node: Some( + ResTarget( + ResTarget { + name: "", + indirection: [], + val: Some( + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 26, + val: Some( + Ival( + Integer { + ival: 2, + }, + ), + ), + }, + ), + ), + }, + ), + location: 26, + }, + ), + ), + }, + ], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopNone, + all: false, + larg: None, + rarg: None, + }, + ), + }, + ), + range: 0..27, + }, + ], +} diff --git a/crates/parser/tests/snapshots/statements/valid/0055@2.snap b/crates/parser/tests/snapshots/statements/valid/0055@2.snap new file mode 100644 index 00000000..96be67fc --- /dev/null +++ b/crates/parser/tests/snapshots/statements/valid/0055@2.snap @@ -0,0 +1,163 @@ +--- +source: crates/parser/tests/statement_parser_test.rs +description: select 1 union select 2; +--- +Parse { + cst: SourceFile@0..24 + SelectStmt@0..24 + Select@0..6 "select" + Whitespace@6..7 " " + SelectStmt@7..8 + ResTarget@7..8 + AConst@7..8 + Iconst@7..8 "1" + Whitespace@8..9 " " + Union@9..14 "union" + Whitespace@14..15 " " + SelectStmt@15..23 + Select@15..21 "select" + Whitespace@21..22 " " + ResTarget@22..23 + AConst@22..23 + Iconst@22..23 "2" + Ascii59@23..24 ";" + , + errors: [], + stmts: [ + RawStmt { + stmt: SelectStmt( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopUnion, + all: false, + larg: Some( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [ + Node { + node: Some( + ResTarget( + ResTarget { + name: "", + indirection: [], + val: Some( + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 7, + val: Some( + Ival( + Integer { + ival: 1, + }, + ), + ), + }, + ), + ), + }, + ), + location: 7, + }, + ), + ), + }, + ], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopNone, + all: false, + larg: None, + rarg: None, + }, + ), + rarg: Some( + SelectStmt { + distinct_clause: [], + into_clause: None, + target_list: [ + Node { + node: Some( + ResTarget( + ResTarget { + name: "", + indirection: [], + val: Some( + Node { + node: Some( + AConst( + AConst { + isnull: false, + location: 22, + val: Some( + Ival( + Integer { + ival: 2, + }, + ), + ), + }, + ), + ), + }, + ), + location: 22, + }, + ), + ), + }, + ], + from_clause: [], + where_clause: None, + group_clause: [], + group_distinct: false, + having_clause: None, + window_clause: [], + values_lists: [], + sort_clause: [], + limit_offset: None, + limit_count: None, + limit_option: Default, + locking_clause: [], + with_clause: None, + op: SetopNone, + all: false, + larg: None, + rarg: None, + }, + ), + }, + ), + range: 0..23, + }, + ], +}