progress

Author: psteinroe

crates/pgt_statement_splitter/src/lib.rs

@@ -92,7 +92,7 @@ mod tests {
             assert_eq!(
                 self.result.ranges.len(),
                 expected.len(),
-                "Expected {} statements for input {}, got {}: {:?}",
+                "Expected {} statements for input\n{}\ngot {}:\n{:?}",
                 expected.len(),
                 self.input,
                 self.result.ranges.len(),

crates/pgt_statement_splitter/src/splitter/common.rs

@@ -120,11 +120,17 @@ pub(crate) fn case(p: &mut Splitter) {
 }
 
 pub(crate) fn unknown(p: &mut Splitter, exclude: &[SyntaxKind]) {
+    let mut in_atomic = false;
     loop {
         match p.current() {
             SyntaxKind::SEMICOLON => {
+                if in_atomic {
+                    // only end the statement if the next non-trivia token is not END
+                    // this is to handle cases like BEGIN ATOMIC SELECT ...; END;
+                    p.advance();
+                    break;
+                }
                 p.advance();
-                break;
             }
             SyntaxKind::EOF => {
                 break;
@@ -257,7 +263,6 @@ pub(crate) fn unknown(p: &mut Splitter, exclude: &[SyntaxKind]) {
                     }
                     p.advance();
                 }
-
                 Some(SyntaxKind::CREATE_KW) => {
                     let prev = p.look_back(true);
                     if [
@@ -275,6 +280,17 @@ pub(crate) fn unknown(p: &mut Splitter, exclude: &[SyntaxKind]) {
 
                     p.advance();
                 }
+                Some(SyntaxKind::ATOMIC_KW) => {
+                    if p.look_back(true) == Some(SyntaxKind::BEGIN_KW) {
+                        // BEGIN ATOMIC ... END;
+                        in_atomic = true;
+                    }
+                    p.advance();
+                }
+                Some(SyntaxKind::END_KW) => {
+                    in_atomic = false;
+                    p.advance();
+                }
                 Some(_) => {
                     break;
                 }