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)> = 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/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 { match self { StatementId::Root(_) => None,