feat(completions): insert schema name when selecting non-public tables

Author: juleswritescode

crates/pgt_completions/src/builder.rs

@@ -1,5 +1,5 @@
 use crate::{
-    CompletionItemKind,
+    CompletionItemKind, CompletionText,
     context::CompletionContext,
     item::CompletionItem,
     relevance::{filtering::CompletionFilter, scoring::CompletionScore},
@@ -11,6 +11,7 @@ pub(crate) struct PossibleCompletionItem<'a> {
     pub kind: CompletionItemKind,
     pub score: CompletionScore<'a>,
     pub filter: CompletionFilter<'a>,
+    pub completion_text: Option<CompletionText>,
 }
 
 pub(crate) struct CompletionBuilder<'a> {
@@ -72,6 +73,7 @@ impl<'a> CompletionBuilder<'a> {
 
                     // wonderous Rust syntax ftw
                     sort_text: format!("{:0>padding$}", idx, padding = max_padding),
+                    completion_text: item.completion_text,
                 }
             })
             .collect()

crates/pgt_completions/src/context.rs

@@ -209,14 +209,7 @@ impl<'a> CompletionContext<'a> {
 
         // We have arrived at the leaf node
         if current_node.child_count() == 0 {
-            if matches!(
-                self.get_ts_node_content(current_node).unwrap(),
-                NodeText::Replaced
-            ) {
-                self.node_under_cursor = None;
-            } else {
-                self.node_under_cursor = Some(current_node);
-            }
+            self.node_under_cursor = Some(current_node);
             return;
         }
 

crates/pgt_completions/src/item.rs

@@ -1,5 +1,6 @@
 use std::fmt::Display;
 
+use pgt_text_size::TextRange;
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -25,6 +26,21 @@ impl Display for CompletionItemKind {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
+/// The text that the editor should fill in.
+/// If `None`, the `label` should be used.
+/// Tables, for example, might have different completion_texts:
+///
+/// label: "users", description: "Schema: auth", completion_text: "auth.users".
+pub struct CompletionText {
+    pub text: String,
+    /// A `range` is required because some editors replace the current token,
+    /// others naively insert the text.
+    /// Having a range where start == end makes it an insertion.
+    pub range: TextRange,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
 pub struct CompletionItem {
@@ -34,4 +50,6 @@ pub struct CompletionItem {
     pub kind: CompletionItemKind,
     /// String used for sorting by LSP clients.
     pub sort_text: String,
+
+    pub completion_text: Option<CompletionText>,
 }

crates/pgt_completions/src/providers/columns.rs

@@ -17,6 +17,7 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
             filter: CompletionFilter::from(relevance),
             description: format!("Table: {}.{}", col.schema_name, col.table_name),
             kind: CompletionItemKind::Column,
+            completion_text: None,
         };
 
         builder.add_item(item);

crates/pgt_completions/src/providers/functions.rs

@@ -5,6 +5,8 @@ use crate::{
     relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
+use super::helper::get_completion_text_with_schema;
+
 pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
     let available_functions = &ctx.schema_cache.functions;
 
@@ -17,6 +19,7 @@ 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(ctx, &func.name, &func.schema),
         };
 
         builder.add_item(item);

crates/pgt_completions/src/providers/helper.rs

@@ -0,0 +1,27 @@
+use pgt_text_size::{TextRange, TextSize};
+
+use crate::{CompletionText, context::CompletionContext};
+
+pub(crate) fn get_completion_text_with_schema(
+    ctx: &CompletionContext,
+    item_name: &str,
+    item_schema_name: &str,
+) -> Option<CompletionText> {
+    if item_schema_name == "public" {
+        None
+    } else if ctx.schema_name.is_some() {
+        None
+    } else {
+        let node = ctx.node_under_cursor.unwrap();
+
+        let range = TextRange::new(
+            TextSize::try_from(node.start_byte()).unwrap(),
+            TextSize::try_from(node.end_byte()).unwrap(),
+        );
+
+        Some(CompletionText {
+            text: format!("{}.{}", item_schema_name, item_name),
+            range,
+        })
+    }
+}

crates/pgt_completions/src/providers/mod.rs

@@ -1,5 +1,6 @@
 mod columns;
 mod functions;
+mod helper;
 mod schemas;
 mod tables;
 

crates/pgt_completions/src/providers/schemas.rs

@@ -16,6 +16,7 @@ pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut Completion
             kind: crate::CompletionItemKind::Schema,
             score: CompletionScore::from(relevance.clone()),
             filter: CompletionFilter::from(relevance),
+            completion_text: None,
         };
 
         builder.add_item(item);

crates/pgt_completions/src/providers/tables.rs

@@ -5,6 +5,8 @@ use crate::{
     relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
+use super::helper::get_completion_text_with_schema;
+
 pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
     let available_tables = &ctx.schema_cache.tables;
 
@@ -17,6 +19,7 @@ pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionB
             filter: CompletionFilter::from(relevance),
             description: format!("Schema: {}", table.schema),
             kind: CompletionItemKind::Table,
+            completion_text: get_completion_text_with_schema(ctx, &table.name, &table.schema),
         };
 
         builder.add_item(item);

crates/pgt_lsp/src/handlers/completions.rs

@@ -1,16 +1,23 @@
-use crate::{adapters::get_cursor_position, session::Session};
+use crate::{
+    adapters::{self, get_cursor_position},
+    diagnostics::LspError,
+    session::Session,
+};
 use anyhow::Result;
 use pgt_workspace::{WorkspaceError, features::completions::GetCompletionsParams};
-use tower_lsp::lsp_types::{self, CompletionItem, CompletionItemLabelDetails};
+use tower_lsp::lsp_types::{self, CompletionItem, CompletionItemLabelDetails, TextEdit};
 
 #[tracing::instrument(level = "debug", skip(session), err)]
 pub fn get_completions(
     session: &Session,
     params: lsp_types::CompletionParams,
-) -> Result<lsp_types::CompletionResponse> {
+) -> Result<lsp_types::CompletionResponse, LspError> {
     let url = params.text_document_position.text_document.uri;
     let path = session.file_path(&url)?;
 
+    let doc = session.document(&url)?;
+    let encoding = adapters::negotiated_encoding(session.client_capabilities().unwrap());
+
     let completion_result = match session.workspace.get_completions(GetCompletionsParams {
         path,
         position: get_cursor_position(session, &url, params.text_document_position.position)?,
@@ -36,6 +43,12 @@ pub fn get_completions(
             }),
             preselect: Some(i.preselected),
             sort_text: Some(i.sort_text),
+            text_edit: i.completion_text.map(|c| {
+                lsp_types::CompletionTextEdit::Edit(TextEdit {
+                    new_text: c.text,
+                    range: adapters::to_lsp::range(&doc.line_index, c.range, encoding).unwrap(),
+                })
+            }),
             kind: Some(to_lsp_types_completion_item_kind(i.kind)),
             ..CompletionItem::default()
         })