fix: begin atomic split

Author: psteinroe

crates/pgt_statement_splitter/src/lib.rs

@@ -133,6 +133,30 @@ mod tests {
         }
     }
 
+    #[test]
+    fn begin_atomic() {
+        Tester::from(
+            "CREATE OR REPLACE FUNCTION public.test_fn(some_in TEXT)
+RETURNS TEXT
+LANGUAGE sql
+IMMUTABLE
+STRICT
+BEGIN ATOMIC
+  SELECT $1 || 'foo';
+END;",
+        )
+        .expect_statements(vec![
+            "CREATE OR REPLACE FUNCTION public.test_fn(some_in TEXT)
+RETURNS TEXT
+LANGUAGE sql
+IMMUTABLE
+STRICT
+BEGIN ATOMIC
+  SELECT $1 || 'foo';
+END;",
+        ]);
+    }
+
     #[test]
     fn ts_with_timezone() {
         Tester::from("alter table foo add column bar timestamp with time zone;").expect_statements(

crates/pgt_statement_splitter/src/splitter/common.rs

@@ -188,6 +188,8 @@ pub(crate) fn unknown(p: &mut Splitter, exclude: &[SyntaxKind]) {
                         // for revoke
                         SyntaxKind::REVOKE_KW,
                         SyntaxKind::COMMA,
+                        // for BEGIN ATOMIC
+                        SyntaxKind::ATOMIC_KW,
                     ]
                     .iter()
                     .all(|x| Some(x) != prev.as_ref())