feat(completions): nicer function completions (#401)

Author: juleswritescode

crates/pgt_completions/src/item.rs

@@ -41,6 +41,8 @@ pub struct CompletionText {
     /// others naively insert the text.
     /// Having a range where start == end makes it an insertion.
     pub range: TextRange,
+
+    pub is_snippet: bool,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

crates/pgt_completions/src/providers/functions.rs

@@ -1,7 +1,10 @@
+use pgt_schema_cache::Function;
+
 use crate::{
-    CompletionItemKind,
+    CompletionItemKind, CompletionText,
     builder::{CompletionBuilder, PossibleCompletionItem},
     context::CompletionContext,
+    providers::helper::get_range_to_replace,
     relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
@@ -19,17 +22,46 @@ pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut Completi
             filter: CompletionFilter::from(relevance),
             description: format!("Schema: {}", func.schema),
             kind: CompletionItemKind::Function,
-            completion_text: get_completion_text_with_schema_or_alias(
-                ctx,
-                &func.name,
-                &func.schema,
-            ),
+            completion_text: Some(get_completion_text(ctx, func)),
         };
 
         builder.add_item(item);
     }
 }
 
+fn get_completion_text(ctx: &CompletionContext, func: &Function) -> CompletionText {
+    let range = get_range_to_replace(ctx);
+    let mut text = get_completion_text_with_schema_or_alias(ctx, &func.name, &func.schema)
+        .map(|ct| ct.text)
+        .unwrap_or(func.name.to_string());
+
+    if ctx.is_invocation {
+        CompletionText {
+            text,
+            range,
+            is_snippet: false,
+        }
+    } else {
+        text.push('(');
+
+        let num_args = func.args.args.len();
+        for (idx, arg) in func.args.args.iter().enumerate() {
+            text.push_str(format!(r#"${{{}:{}}}"#, idx + 1, arg.name).as_str());
+            if idx < num_args - 1 {
+                text.push_str(", ");
+            }
+        }
+
+        text.push(')');
+
+        CompletionText {
+            text,
+            range,
+            is_snippet: num_args > 0,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{

crates/pgt_completions/src/providers/helper.rs

@@ -34,14 +34,17 @@ pub(crate) fn get_completion_text_with_schema_or_alias(
     item_name: &str,
     schema_or_alias_name: &str,
 ) -> Option<CompletionText> {
-    if schema_or_alias_name == "public" || ctx.schema_or_alias_name.is_some() {
+    let is_already_prefixed_with_schema_name = ctx.schema_or_alias_name.is_some();
+
+    if schema_or_alias_name == "public" || is_already_prefixed_with_schema_name {
         None
     } else {
         let range = get_range_to_replace(ctx);
 
         Some(CompletionText {
             text: format!("{}.{}", schema_or_alias_name, item_name),
             range,
+            is_snippet: false,
         })
     }
 }

crates/pgt_completions/src/providers/policies.rs

@@ -25,6 +25,7 @@ pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut Completi
             let range = get_range_to_replace(ctx);
             Some(CompletionText {
                 text: pol.name.clone(),
+                is_snippet: false,
                 range: TextRange::new(
                     range.start() + TextSize::new(1),
                     range.end() - TextSize::new(1),
@@ -34,6 +35,7 @@ pub fn complete_policies<'a>(ctx: &CompletionContext<'a>, builder: &mut Completi
             // If we aren't within quotes, we want to complete the
             // full policy including quotation marks.
             Some(CompletionText {
+                is_snippet: false,
                 text: format!("\"{}\"", pol.name),
                 range: get_range_to_replace(ctx),
             })

crates/pgt_lsp/src/handlers/completions.rs

@@ -5,7 +5,9 @@ use crate::{
 };
 use anyhow::Result;
 use pgt_workspace::{WorkspaceError, features::completions::GetCompletionsParams};
-use tower_lsp::lsp_types::{self, CompletionItem, CompletionItemLabelDetails, TextEdit};
+use tower_lsp::lsp_types::{
+    self, CompletionItem, CompletionItemLabelDetails, InsertTextFormat, TextEdit,
+};
 
 #[tracing::instrument(level = "debug", skip(session), err)]
 pub fn get_completions(
@@ -43,6 +45,11 @@ pub fn get_completions(
             }),
             preselect: Some(i.preselected),
             sort_text: Some(i.sort_text),
+            insert_text_format: if i.completion_text.as_ref().is_some_and(|c| c.is_snippet) {
+                Some(InsertTextFormat::SNIPPET)
+            } else {
+                Some(InsertTextFormat::PLAIN_TEXT)
+            },
             text_edit: i.completion_text.map(|c| {
                 lsp_types::CompletionTextEdit::Edit(TextEdit {
                     new_text: c.text,

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

@@ -15,6 +15,7 @@ pub struct SQLFunctionArg {
 
 #[derive(Debug, Clone)]
 pub struct SQLFunctionSignature {
+    #[allow(dead_code)]
     pub schema: Option<String>,
     pub name: String,
     pub args: Vec<SQLFunctionArg>,