merge with main

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;

crates/pgt_workspace/src/features/completions.rs

@@ -49,18 +49,24 @@ pub(crate) fn get_statement_for_completions(
     if count == 1 {
         eligible_statements.next()
     } else {
-        let mut prev_stmt = None;
+        let mut prev_stmt: Option<(StatementId, TextRange, String, Arc<tree_sitter::Tree>)> = None;
 
         for current_stmt in eligible_statements {
             /*
              * If we have multiple statements, we want to make sure that we do not overlap
              * with the next one.
              *
              * select 1 |select 1;
+             *
+             * This is however ok if the current statement is a child of the previous one,
+             * such as in CREATE FUNCTION bodies.
              */
-            if prev_stmt.is_some_and(|_| current_stmt.1.contains(position)) {
+            if prev_stmt.is_some_and(|prev| {
+                current_stmt.1.contains(position) && !current_stmt.0.is_child_of(&prev.0)
+            }) {
                 return None;
             }
+
             prev_stmt = Some(current_stmt)
         }
 
@@ -162,6 +168,30 @@ mod tests {
         assert_eq!(text, "select * from")
     }
 
+    #[test]
+    fn identifies_nested_stmts() {
+        let sql = format!(
+            r#"
+            create or replace function one()
+            returns integer
+            language sql
+            as $$
+                select {} from cool;
+            $$;
+        "#,
+            CURSOR_POSITION
+        );
+
+        let sql = sql.trim();
+
+        let (doc, position) = get_doc_and_pos(sql);
+
+        let (_, _, text, _) =
+            get_statement_for_completions(&doc, position).expect("Expected Statement");
+
+        assert_eq!(text.trim(), "select  from cool;")
+    }
+
     #[test]
     fn does_not_consider_too_far_offset() {
         let sql = format!("select * from  {}", CURSOR_POSITION);

crates/pgt_workspace/src/workspace/server/statement_identifier.rs

@@ -66,6 +66,17 @@ impl StatementId {
         matches!(self, StatementId::Child(_))
     }
 
+    pub fn is_child_of(&self, maybe_parent: &StatementId) -> bool {
+        match self {
+            StatementId::Root(_) => false,
+            StatementId::Child(child_root) => match maybe_parent {
+                StatementId::Root(parent_rood) => child_root == parent_rood,
+                // TODO: can we have multiple nested statements?
+                StatementId::Child(_) => false,
+            },
+        }
+    }
+
     pub fn parent(&self) -> Option<StatementId> {
         match self {
             StatementId::Root(_) => None,