diff --git a/crates/pgt_statement_splitter/src/lib.rs b/crates/pgt_statement_splitter/src/lib.rs index e43a1095..c53ae78c 100644 --- a/crates/pgt_statement_splitter/src/lib.rs +++ b/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] diff --git a/crates/pgt_statement_splitter/src/parser/common.rs b/crates/pgt_statement_splitter/src/parser/common.rs index a5d68df1..2498c04a 100644 --- a/crates/pgt_statement_splitter/src/parser/common.rs +++ b/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; } diff --git a/crates/pgt_statement_splitter/tests/data/grant_statements__4.sql b/crates/pgt_statement_splitter/tests/data/grant_statements__4.sql new file mode 100644 index 00000000..5811810e --- /dev/null +++ b/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;