feat(compl): complete in (simple) function bodies (#426)

Author: juleswritescode

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,