From d282e89b118452cd2f1ef284217bd5b2f263326e Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Sun, 18 Jan 2026 12:20:30 -0500 Subject: [PATCH] ide: delete from completion improvements --- crates/squawk_ide/src/binder.rs | 22 ++ crates/squawk_ide/src/completion.rs | 403 ++++++++++++++++++++++---- crates/squawk_ide/src/resolve.rs | 4 +- crates/squawk_server/src/lsp_utils.rs | 18 +- 4 files changed, 374 insertions(+), 73 deletions(-) diff --git a/crates/squawk_ide/src/binder.rs b/crates/squawk_ide/src/binder.rs index d0fb6a0f..4adcd6bb 100644 --- a/crates/squawk_ide/src/binder.rs +++ b/crates/squawk_ide/src/binder.rs @@ -204,6 +204,28 @@ impl Binder { } names } + + pub(crate) fn functions_with_single_param(&self, param_type: &Name) -> Vec<&Name> { + let root_scope = self.root_scope(); + let scope = &self.scopes[root_scope]; + + let mut names = vec![]; + for (name, symbol_ids) in &scope.entries { + for symbol_id in symbol_ids { + let symbol = &self.symbols[*symbol_id]; + if symbol.kind == SymbolKind::Function + && symbol + .params + .as_ref() + .is_some_and(|p| p.len() == 1 && &p[0] == param_type) + { + names.push(name); + break; + } + } + } + names + } } pub(crate) fn bind(file: &ast::SourceFile) -> Binder { diff --git a/crates/squawk_ide/src/completion.rs b/crates/squawk_ide/src/completion.rs index cb0145bb..12df6f36 100644 --- a/crates/squawk_ide/src/completion.rs +++ b/crates/squawk_ide/src/completion.rs @@ -1,7 +1,6 @@ use rowan::TextSize; use squawk_syntax::ast::{self, AstNode}; use squawk_syntax::{SyntaxKind, SyntaxToken}; -use std::collections::HashSet; use crate::binder; use crate::resolve; @@ -26,6 +25,8 @@ pub fn completion(file: &ast::SourceFile, offset: TextSize) -> Vec { select_completions(file, select_clause, &token) } + CompletionContext::DeleteClauses(delete) => delete_clauses_completions(&delete), + CompletionContext::DeleteExpr(delete) => delete_expr_completions(file, &delete, &token), } } @@ -45,6 +46,7 @@ fn select_completions( insert_text: None, insert_text_format: None, trigger_completion_after_insert: false, + sort_text: None, })); let tables = binder.all_symbols_by_kind(SymbolKind::Table, schema.as_ref()); @@ -55,6 +57,7 @@ fn select_completions( insert_text: None, insert_text_format: None, trigger_completion_after_insert: false, + sort_text: None, })); if schema.is_none() { @@ -81,6 +84,7 @@ fn select_completions( insert_text: None, insert_text_format: None, trigger_completion_after_insert: false, + sort_text: None, }) })); } @@ -91,30 +95,40 @@ fn select_completions( } fn schema_completions(binder: &binder::Binder) -> Vec { - let default_schemas = - ["public", "pg_temp", "pg_catalog", "pg_toast", "postgres"].map(Name::from_string); - let mut names = HashSet::from(default_schemas); - for name in binder.all_symbols_by_kind(SymbolKind::Schema, None) { - names.insert(name.clone()); - } - - names + let builtin_schemas = ["public", "pg_catalog", "pg_temp", "pg_toast", "postgres"]; + let mut completions: Vec = builtin_schemas .into_iter() - .map(|name| CompletionItem { + .enumerate() + .map(|(i, name)| CompletionItem { label: name.to_string(), kind: CompletionItemKind::Schema, detail: None, insert_text: None, insert_text_format: None, trigger_completion_after_insert: false, + sort_text: Some(format!("{i}")), }) - .collect() + .collect(); + + for name in binder.all_symbols_by_kind(SymbolKind::Schema, None) { + completions.push(CompletionItem { + label: name.to_string(), + kind: CompletionItemKind::Schema, + detail: None, + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: None, + }); + } + + completions } fn table_completions(file: &ast::SourceFile, token: &SyntaxToken) -> Vec { let binder = binder::bind(file); - // We're in a TRUNCATE or TABLE statement, return table names - let tables = binder.all_symbols_by_kind(SymbolKind::Table, None); + let schema = schema_qualifier_at_token(token); + let tables = binder.all_symbols_by_kind(SymbolKind::Table, schema.as_ref()); let mut completions: Vec = tables .into_iter() .map(|name| CompletionItem { @@ -124,28 +138,184 @@ fn table_completions(file: &ast::SourceFile, token: &SyntaxToken) -> Vec Vec { + let mut completions = vec![]; + + if delete.using_clause().is_none() { + completions.push(CompletionItem { + label: "using".to_owned(), + kind: CompletionItemKind::Keyword, + detail: None, + insert_text: Some("using $0".to_owned()), + insert_text_format: Some(CompletionInsertTextFormat::Snippet), + trigger_completion_after_insert: true, + sort_text: None, + }); + } + + if delete.where_clause().is_none() { + completions.push(CompletionItem { + label: "where".to_owned(), + kind: CompletionItemKind::Keyword, + detail: None, + insert_text: Some("where $0".to_owned()), + insert_text_format: Some(CompletionInsertTextFormat::Snippet), + trigger_completion_after_insert: true, + sort_text: None, + }); + } + + if delete.returning_clause().is_none() { + completions.push(CompletionItem { + label: "returning".to_owned(), + kind: CompletionItemKind::Keyword, + detail: None, + insert_text: Some("returning $0".to_owned()), + insert_text_format: Some(CompletionInsertTextFormat::Snippet), + trigger_completion_after_insert: true, + sort_text: None, + }); + } + + completions +} + +fn delete_expr_completions( + file: &ast::SourceFile, + delete: &ast::Delete, + token: &SyntaxToken, +) -> Vec { + let binder = binder::bind(file); + let mut completions = vec![]; + + let Some(path) = delete.relation_name().and_then(|r| r.path()) else { + return completions; + }; + let Some(delete_table_name) = resolve::extract_table_name(&path) else { + return completions; + }; + + let has_table_qualifier = qualifier_at_token(token).is_some_and(|q| q == delete_table_name); + + if has_table_qualifier { + let functions = binder.functions_with_single_param(&delete_table_name); + completions.extend(functions.into_iter().map(|name| CompletionItem { + label: name.to_string(), + kind: CompletionItemKind::Function, + detail: None, + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: None, + })); + } else { + let functions = binder.all_symbols_by_kind(SymbolKind::Function, None); + completions.extend(functions.into_iter().map(|name| CompletionItem { + label: format!("{name}()"), + kind: CompletionItemKind::Function, + detail: None, + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: None, + })); + + completions.push(CompletionItem { + label: delete_table_name.to_string(), + kind: CompletionItemKind::Table, + detail: None, + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: None, + }); + } + + let schema = resolve::extract_schema_name(&path); + let position = path.syntax().text_range().start(); + if let Some(table_ptr) = + binder.lookup_with(&delete_table_name, SymbolKind::Table, position, &schema) + && let Some(create_table) = table_ptr + .to_node(file.syntax()) + .ancestors() + .find_map(ast::CreateTableLike::cast) + { + let columns = resolve::collect_table_columns(&binder, file.syntax(), &create_table); + completions.extend(columns.into_iter().filter_map(|column| { + let name = column.name()?; + Some(CompletionItem { + label: Name::from_node(&name).to_string(), + kind: CompletionItemKind::Column, + detail: None, + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: None, + }) + })); + } + + completions +} + +fn qualifier_at_token(token: &SyntaxToken) -> Option { + let qualifier_token = if token.kind() == SyntaxKind::DOT { + token.prev_token() + } else if token.kind() == SyntaxKind::IDENT + && let Some(prev) = token.prev_token() + && prev.kind() == SyntaxKind::DOT + { + prev.prev_token() + } else { + None + }; + + qualifier_token + .filter(|tk| tk.kind() == SyntaxKind::IDENT) + .map(|tk| Name::from_string(tk.text().to_string())) +} + enum CompletionContext { TableOnly, Default, SelectClause(ast::SelectClause), + DeleteClauses(ast::Delete), + DeleteExpr(ast::Delete), } fn completion_context(token: &SyntaxToken) -> CompletionContext { if let Some(node) = token.parent() { + let mut inside_delete_clause = false; for a in node.ancestors() { if ast::Truncate::can_cast(a.kind()) || ast::Table::can_cast(a.kind()) { return CompletionContext::TableOnly; } + if ast::WhereClause::can_cast(a.kind()) + || ast::UsingClause::can_cast(a.kind()) + || ast::ReturningClause::can_cast(a.kind()) + { + inside_delete_clause = true; + } + if let Some(delete) = ast::Delete::cast(a.clone()) { + if inside_delete_clause { + return CompletionContext::DeleteExpr(delete); + } + if delete.relation_name().is_some() { + return CompletionContext::DeleteClauses(delete); + } + return CompletionContext::TableOnly; + } if let Some(select_clause) = ast::SelectClause::cast(a.clone()) { return CompletionContext::SelectClause(select_clause); } @@ -168,24 +338,11 @@ fn token_at_offset(file: &ast::SourceFile, offset: TextSize) -> Option Option { - let schema_token = if token.kind() == SyntaxKind::DOT { - token.prev_token() - } else if token.kind() == SyntaxKind::IDENT - && let Some(prev) = token.prev_token() - && prev.kind() == SyntaxKind::DOT - { - prev.prev_token() - } else { - None - }; - - schema_token - .filter(|tk| tk.kind() == SyntaxKind::IDENT) - .map(|tk| Schema(Name::from_string(tk.text().to_string()))) + qualifier_at_token(token).map(Schema) } fn default_completions() -> Vec { - ["select", "table", "truncate"] + ["delete from", "select", "table", "truncate"] .map(|stmt| CompletionItem { label: stmt.to_owned(), kind: CompletionItemKind::Keyword, @@ -193,6 +350,7 @@ fn default_completions() -> Vec { insert_text: Some(format!("{stmt} $0;")), insert_text_format: Some(CompletionInsertTextFormat::Snippet), trigger_completion_after_insert: true, + sort_text: None, }) .into_iter() .collect() @@ -209,6 +367,28 @@ pub enum CompletionItemKind { Snippet, } +impl CompletionItemKind { + fn sort_prefix(self) -> &'static str { + match self { + Self::Column => "0", + Self::Keyword => "1", + Self::Table => "1", + Self::Type => "1", + Self::Snippet => "1", + Self::Function => "2", + Self::Schema => "9", + } + } +} + +impl CompletionItem { + pub fn sort_text(&self) -> String { + let prefix = self.kind.sort_prefix(); + let suffix = self.sort_text.as_ref().unwrap_or(&self.label); + format!("{prefix}_{suffix}") + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CompletionInsertTextFormat { PlainText, @@ -223,6 +403,7 @@ pub struct CompletionItem { pub insert_text: Option, pub insert_text_format: Option, pub trigger_completion_after_insert: bool, + pub sort_text: Option, } #[cfg(test)] @@ -258,8 +439,10 @@ mod tests { ) } - fn format_items(items: Vec) -> String { - let mut rows: Vec> = items + fn format_items(mut items: Vec) -> String { + items.sort_by_key(|a| a.sort_text()); + + let rows: Vec> = items .into_iter() .map(|item| { vec![ @@ -271,8 +454,6 @@ mod tests { }) .collect(); - rows.sort(); - let mut builder = Builder::default(); builder.push_record(["label", "kind", "detail", "insert_text"]); for row in rows { @@ -287,11 +468,12 @@ mod tests { #[test] fn completion_at_start() { assert_snapshot!(completions("$0"), @r" - label | kind | detail | insert_text - ----------+---------+--------+-------------- - select | Keyword | | select $0; - table | Keyword | | table $0; - truncate | Keyword | | truncate $0; + label | kind | detail | insert_text + -------------+---------+--------+----------------- + delete from | Keyword | | delete from $0; + select | Keyword | | select $0; + table | Keyword | | table $0; + truncate | Keyword | | truncate $0; "); } @@ -301,11 +483,12 @@ mod tests { create table t(a int); $0 "), @r" - label | kind | detail | insert_text - ----------+---------+--------+-------------- - select | Keyword | | select $0; - table | Keyword | | table $0; - truncate | Keyword | | truncate $0; + label | kind | detail | insert_text + -------------+---------+--------+----------------- + delete from | Keyword | | delete from $0; + select | Keyword | | select $0; + table | Keyword | | table $0; + truncate | Keyword | | truncate $0; "); } @@ -327,34 +510,36 @@ truncate $0; "), @r" label | kind | detail | insert_text ------------+--------+--------+------------- + users | Table | | + public | Schema | | pg_catalog | Schema | | pg_temp | Schema | | pg_toast | Schema | | - postgres | Schema | | - public | Schema | | - users | Table | | + postgres | Schema | | "); } #[test] fn completion_table_at_top_level() { assert_snapshot!(completions("$0"), @r" - label | kind | detail | insert_text - ----------+---------+--------+-------------- - select | Keyword | | select $0; - table | Keyword | | table $0; - truncate | Keyword | | truncate $0; + label | kind | detail | insert_text + -------------+---------+--------+----------------- + delete from | Keyword | | delete from $0; + select | Keyword | | select $0; + table | Keyword | | table $0; + truncate | Keyword | | truncate $0; "); } #[test] fn completion_table_nested() { assert_snapshot!(completions("select * from ($0)"), @r" - label | kind | detail | insert_text - ----------+---------+--------+-------------- - select | Keyword | | select $0; - table | Keyword | | table $0; - truncate | Keyword | | truncate $0; + label | kind | detail | insert_text + -------------+---------+--------+----------------- + delete from | Keyword | | delete from $0; + select | Keyword | | select $0; + table | Keyword | | table $0; + truncate | Keyword | | truncate $0; "); } @@ -366,12 +551,12 @@ table $0; "), @r" label | kind | detail | insert_text ------------+--------+--------+------------- + users | Table | | + public | Schema | | pg_catalog | Schema | | pg_temp | Schema | | pg_toast | Schema | | - postgres | Schema | | - public | Schema | | - users | Table | | + postgres | Schema | | "); } @@ -386,13 +571,13 @@ select $0 from t; ------------+----------+--------+------------- a | Column | | b | Column | | + t | Table | | f() | Function | | + public | Schema | | pg_catalog | Schema | | pg_temp | Schema | | pg_toast | Schema | | - postgres | Schema | | - public | Schema | | - t | Table | | + postgres | Schema | | "); } @@ -408,4 +593,98 @@ select public.$0; f() | Function | | "); } + + #[test] + fn completion_truncate_with_schema_qualifier() { + assert_snapshot!(completions(" +create table users (id int); +truncate public.$0; +"), @r" + label | kind | detail | insert_text + -------+-------+--------+------------- + users | Table | | + "); + } + + #[test] + fn completion_after_delete_from() { + assert_snapshot!(completions(" +create table users (id int); +delete from $0; +"), @r" + label | kind | detail | insert_text + ------------+--------+--------+------------- + users | Table | | + public | Schema | | + pg_catalog | Schema | | + pg_temp | Schema | | + pg_toast | Schema | | + postgres | Schema | | + "); + } + + #[test] + fn completion_delete_clauses() { + assert_snapshot!(completions(" +create table t (id int); +delete from t $0; +"), @r" + label | kind | detail | insert_text + -----------+---------+--------+-------------- + returning | Keyword | | returning $0 + using | Keyword | | using $0 + where | Keyword | | where $0 + "); + } + + #[test] + fn completion_delete_where_expr() { + assert_snapshot!(completions(" +create table t (id int, name text); +create function is_active() returns bool as 'select true' language sql; +delete from t where $0; +"), @r" + label | kind | detail | insert_text + -------------+----------+--------+------------- + id | Column | | + name | Column | | + t | Table | | + is_active() | Function | | + ") + } + + #[test] + fn completion_delete_returning_expr() { + assert_snapshot!(completions(" +create table t (id int, name text); +delete from t returning $0; +"), @r" + label | kind | detail | insert_text + -------+--------+--------+------------- + id | Column | | + name | Column | | + t | Table | | + "); + } + + #[test] + fn completion_delete_where_qualified() { + assert_snapshot!(completions(" +-- different type than the table, so we shouldn't show this +create function b(diff_type) returns int8 + as 'select 1' + language sql; +create function f(t) returns int8 + as 'select 1' + language sql; +create table t (a int, b text); +delete from t where t.$0; +"), @r" + label | kind | detail | insert_text + -------+----------+--------+------------- + a | Column | | + b | Column | | + f | Function | | + "); + } } diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 9cc5e6c0..465f040f 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1598,13 +1598,13 @@ fn find_containing_path(name_ref: &ast::NameRef) -> Option { None } -fn extract_table_name(path: &ast::Path) -> Option { +pub(crate) fn extract_table_name(path: &ast::Path) -> Option { let segment = path.segment()?; let name_ref = segment.name_ref()?; Some(Name::from_node(&name_ref)) } -fn extract_schema_name(path: &ast::Path) -> Option { +pub(crate) fn extract_schema_name(path: &ast::Path) -> Option { path.qualifier() .and_then(|q| q.segment()) .and_then(|s| s.name_ref()) diff --git a/crates/squawk_server/src/lsp_utils.rs b/crates/squawk_server/src/lsp_utils.rs index bc949ec6..0d1bf95f 100644 --- a/crates/squawk_server/src/lsp_utils.rs +++ b/crates/squawk_server/src/lsp_utils.rs @@ -80,17 +80,17 @@ pub(crate) fn completion_item( ) -> lsp_types::CompletionItem { use squawk_ide::completion::{CompletionInsertTextFormat, CompletionItemKind}; - let (sort_prefix, kind) = match item.kind { - CompletionItemKind::Schema => ("9", lsp_types::CompletionItemKind::MODULE), - CompletionItemKind::Keyword => ("0", lsp_types::CompletionItemKind::KEYWORD), - CompletionItemKind::Table => ("0", lsp_types::CompletionItemKind::STRUCT), - CompletionItemKind::Column => ("0", lsp_types::CompletionItemKind::FIELD), - CompletionItemKind::Function => ("0", lsp_types::CompletionItemKind::FUNCTION), - CompletionItemKind::Type => ("0", lsp_types::CompletionItemKind::CLASS), - CompletionItemKind::Snippet => ("0", lsp_types::CompletionItemKind::SNIPPET), + let kind = match item.kind { + CompletionItemKind::Schema => lsp_types::CompletionItemKind::MODULE, + CompletionItemKind::Keyword => lsp_types::CompletionItemKind::KEYWORD, + CompletionItemKind::Table => lsp_types::CompletionItemKind::STRUCT, + CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD, + CompletionItemKind::Function => lsp_types::CompletionItemKind::FUNCTION, + CompletionItemKind::Type => lsp_types::CompletionItemKind::CLASS, + CompletionItemKind::Snippet => lsp_types::CompletionItemKind::SNIPPET, }; - let sort_text = Some(format!("{sort_prefix}_{}", item.label)); + let sort_text = Some(item.sort_text()); let insert_text_format = item.insert_text_format.map(|x| match x { CompletionInsertTextFormat::PlainText => lsp_types::InsertTextFormat::PLAIN_TEXT,