fix(splitter): split grant stmts (#425)

Author: juleswritescode

crates/pgt_statement_splitter/src/lib.rs

@@ -45,8 +45,9 @@ mod tests {
             assert_eq!(
                 self.parse.ranges.len(),
                 expected.len(),
-                "Expected {} statements, got {}: {:?}",
+                "Expected {} statements for input {}, got {}: {:?}",
                 expected.len(),
+                self.input,
                 self.parse.ranges.len(),
                 self.parse
                     .ranges
@@ -114,10 +115,24 @@ mod tests {
 
     #[test]
     fn grant() {
-        Tester::from("GRANT SELECT ON TABLE \"public\".\"my_table\" TO \"my_role\";")
-            .expect_statements(vec![
-                "GRANT SELECT ON TABLE \"public\".\"my_table\" TO \"my_role\";",
-            ]);
+        let stmts = vec![
+            "GRANT SELECT ON TABLE \"public\".\"my_table\" TO \"my_role\";",
+            "GRANT UPDATE ON TABLE \"public\".\"my_table\" TO \"my_role\";",
+            "GRANT DELETE ON TABLE \"public\".\"my_table\" TO \"my_role\";",
+            "GRANT INSERT ON TABLE \"public\".\"my_table\" TO \"my_role\";",
+            "GRANT CREATE ON SCHEMA \"public\" TO \"my_role\";",
+            "GRANT ALL PRIVILEGES ON DATABASE \"my_database\" TO \"my_role\";",
+            "GRANT USAGE ON SCHEMA \"public\" TO \"my_role\";",
+            "GRANT EXECUTE ON FUNCTION \"public\".\"my_function\"() TO \"my_role\";",
+            "GRANT REFERENCES ON TABLE \"public\".\"my_table\" TO \"my_role\";",
+            "GRANT SELECT, UPDATE ON ALL TABLES IN SCHEMA \"public\" TO \"my_role\";",
+            "GRANT SELECT, INSERT ON public.users TO anon WITH GRANT OPION GRANTED BY owner;",
+            "GRANT owner, admin to anon WITH ADMIN;",
+        ];
+
+        for stmt in stmts {
+            Tester::from(stmt).expect_statements(vec![stmt]);
+        }
     }
 
     #[test]

crates/pgt_statement_splitter/src/parser/common.rs

@@ -217,6 +217,7 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
                         SyntaxKind::Except,
                         // for grant
                         SyntaxKind::Grant,
+                        SyntaxKind::Ascii44,
                     ]
                     .iter()
                     .all(|x| Some(x) != prev.as_ref())
@@ -246,6 +247,7 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
                         SyntaxKind::Instead,
                         // for grant
                         SyntaxKind::Grant,
+                        SyntaxKind::Ascii44,
                     ]
                     .iter()
                     .all(|x| Some(x) != prev.as_ref())
@@ -263,6 +265,10 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
                         SyntaxKind::Check,
                         // TIMESTAMP WITH TIME ZONE should not start a new statement
                         SyntaxKind::Time,
+                        SyntaxKind::Grant,
+                        SyntaxKind::Admin,
+                        SyntaxKind::Inherit,
+                        SyntaxKind::Set,
                     ]
                     .iter()
                     .all(|x| Some(x) != next.as_ref())
@@ -271,6 +277,22 @@ pub(crate) fn unknown(p: &mut Parser, exclude: &[SyntaxKind]) {
                     }
                     p.advance();
                 }
+
+                Some(SyntaxKind::Create) => {
+                    let prev = p.look_back().map(|t| t.kind);
+                    if [
+                        // for grant
+                        SyntaxKind::Grant,
+                        SyntaxKind::Ascii44,
+                    ]
+                    .iter()
+                    .all(|x| Some(x) != prev.as_ref())
+                    {
+                        break;
+                    }
+
+                    p.advance();
+                }
                 Some(_) => {
                     break;
                 }

crates/pgt_statement_splitter/tests/data/grant_statements__4.sql

@@ -0,0 +1,11 @@
+GRANT CREATE ON SCHEMA public TO anon;
+
+GRANT SELECT, INSERT ON public.users TO anon WITH GRANT OPTION GRANTED BY Owner;
+
+GRANT read_access, write_access TO user_role
+  WITH INHERIT TRUE
+  GRANTED BY security_admin;
+
+GRANT manager_role TO employee_role
+  WITH ADMIN OPTION
+  GRANTED BY admin_role;