From 218fcfd4f2f8f94ded8fd6ed3e64cc04c08a537b Mon Sep 17 00:00:00 2001
From: Julian <juliandomke@outlook.de>
Date: Mon, 9 Jun 2025 12:17:02 +0200
Subject: [PATCH 1/2] ok

---
 .../pgt_workspace/src/features/completions.rs | 34 +++++++++++++++++--
 crates/pgt_workspace/src/workspace/server.rs  |  4 +++
 .../workspace/server/statement_identifier.rs  | 11 ++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/crates/pgt_workspace/src/features/completions.rs b/crates/pgt_workspace/src/features/completions.rs
index 85342183..53eb9eab 100644
--- a/crates/pgt_workspace/src/features/completions.rs
+++ b/crates/pgt_workspace/src/features/completions.rs
@@ -49,7 +49,7 @@ 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 {
             /*
@@ -57,10 +57,16 @@ pub(crate) fn get_statement_for_completions(
              * 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);
diff --git a/crates/pgt_workspace/src/workspace/server.rs b/crates/pgt_workspace/src/workspace/server.rs
index 67a3713c..db4cfdc8 100644
--- a/crates/pgt_workspace/src/workspace/server.rs
+++ b/crates/pgt_workspace/src/workspace/server.rs
@@ -596,12 +596,16 @@ impl Workspace for WorkspaceServer {
 
         let schema_cache = self.schema_cache.load(pool)?;
 
+        tracing::warn!("stmts in doc: {:?}", parsed_doc.iter(DefaultMapper).count());
+
         match get_statement_for_completions(&parsed_doc, params.position) {
             None => {
                 tracing::debug!("No statement found.");
                 Ok(CompletionsResult::default())
             }
             Some((id, range, content, cst)) => {
+                tracing::error!("Found statement with id {:?}: {:#?}", id, content);
+
                 let position = params.position - range.start();
 
                 let items = pgt_completions::complete(pgt_completions::CompletionParams {
diff --git a/crates/pgt_workspace/src/workspace/server/statement_identifier.rs b/crates/pgt_workspace/src/workspace/server/statement_identifier.rs
index 7c7d76f0..627ff261 100644
--- a/crates/pgt_workspace/src/workspace/server/statement_identifier.rs
+++ b/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,

From 1b3f1e0c73b86f1ccf80177a7021ebc2a17ffa0b Mon Sep 17 00:00:00 2001
From: Julian <juliandomke@outlook.de>
Date: Mon, 9 Jun 2025 12:19:16 +0200
Subject: [PATCH 2/2] remove logs

---
 crates/pgt_workspace/src/workspace/server.rs | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/crates/pgt_workspace/src/workspace/server.rs b/crates/pgt_workspace/src/workspace/server.rs
index db4cfdc8..67a3713c 100644
--- a/crates/pgt_workspace/src/workspace/server.rs
+++ b/crates/pgt_workspace/src/workspace/server.rs
@@ -596,16 +596,12 @@ impl Workspace for WorkspaceServer {
 
         let schema_cache = self.schema_cache.load(pool)?;
 
-        tracing::warn!("stmts in doc: {:?}", parsed_doc.iter(DefaultMapper).count());
-
         match get_statement_for_completions(&parsed_doc, params.position) {
             None => {
                 tracing::debug!("No statement found.");
                 Ok(CompletionsResult::default())
             }
             Some((id, range, content, cst)) => {
-                tracing::error!("Found statement with id {:?}: {:#?}", id, content);
-
                 let position = params.position - range.start();
 
                 let items = pgt_completions::complete(pgt_completions::CompletionParams {