fix: parse unions (#329)

Author: juleswritescode

crates/pgt_statement_splitter/src/parser.rs

@@ -147,6 +147,7 @@ impl Parser {
  }
 
  /// Look ahead to the next relevant token
+ /// Returns `None` if we are already at the last relevant token
  fn look_ahead(&self) -> Option<&Token> {
  // we need to look ahead to the next relevant token
  let mut look_ahead_pos = self.next_pos + 1;
@@ -161,6 +162,7 @@ impl Parser {
  }
  }
 
+ /// Returns `None` if there are no previous relevant tokens
  fn look_back(&self) -> Option<&Token> {
  // we need to look back to the last relevant token
  let mut look_back_pos = self.next_pos - 1;

crates/pgt_statement_splitter/src/parser/common.rs

@@ -145,6 +145,12 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
  SyntaxKind::Also,
  // for create rule
  SyntaxKind::Instead,
+ // for UNION
+ SyntaxKind::Union,
+ // for UNION ALL
+ SyntaxKind::All,
+ // for UNION ... EXCEPT
+ SyntaxKind::Except,
  ]
  .iter()
  .all(|x| Some(x) != prev.as_ref())

crates/pgt_statement_splitter/tests/data/simple_union__4.sql

@@ -0,0 +1,7 @@
+select 1 union all select 2;
+
+select 1 union select 2;
+
+select 1 union select 2 except select 3;
+
+select 1 union all select 2 except select 3;