it stopped working, even though I improved it

Author: juleswritescode

crates/pgt_completions/src/providers/columns.rs

@@ -1,13 +1,14 @@
-use pgt_schema_cache::SchemaCache;
-use pgt_treesitter::{TreesitterContext, WrappingClause};
+use pgt_schema_cache::{Column, SchemaCache};
+use pgt_treesitter::TreesitterContext;
 
 use crate::{
- CompletionItemKind,
+ CompletionItemKind, CompletionText,
  builder::{CompletionBuilder, PossibleCompletionItem},
+ providers::helper::{get_range_to_replace, with_closed_quote},
  relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-use super::helper::{find_matching_alias_for_table, get_completion_text_with_schema_or_alias};
+use super::helper::{find_matching_alias_for_table, with_schema_or_alias};
 
 pub fn complete_columns<'a>(
  ctx: &TreesitterContext<'a>,
@@ -19,33 +20,34 @@ pub fn complete_columns<'a>(
  for col in available_columns {
  let relevance = CompletionRelevanceData::Column(col);
 
- let mut item = PossibleCompletionItem {
+ let item = PossibleCompletionItem {
  label: col.name.clone(),
  score: CompletionScore::from(relevance.clone()),
  filter: CompletionFilter::from(relevance),
  description: format!("{}.{}", col.schema_name, col.table_name),
  kind: CompletionItemKind::Column,
- completion_text: None,
+ completion_text: Some(get_completion_text(ctx, col)),
  detail: col.type_name.as_ref().map(|t| t.to_string()),
  };
 
- // autocomplete with the alias in a join clause if we find one
- if matches!(
- ctx.wrapping_clause_type,
- Some(WrappingClause::Join { .. })
- | Some(WrappingClause::Where)
- | Some(WrappingClause::Select)
- ) {
- item.completion_text = find_matching_alias_for_table(ctx, col.table_name.as_str())
- .and_then(|alias| {
- get_completion_text_with_schema_or_alias(ctx, col.name.as_str(), alias.as_str())
- });
- }
-
  builder.add_item(item);
  }
 }
 
+fn get_completion_text(ctx: &TreesitterContext, col: &Column) -> CompletionText {
+ let range = get_range_to_replace(ctx);
+ let col_name = with_closed_quote(ctx, col.name.as_str());
+ let alias = find_matching_alias_for_table(ctx, col.table_name.as_str());
+ let with_schema_or_alias =
+ with_schema_or_alias(ctx, col_name.as_str(), alias.as_ref().map(|s| s.as_str()));
+
+ CompletionText {
+ is_snippet: false,
+ range,
+ text: with_schema_or_alias,
+ }
+}
+
 #[cfg(test)]
 mod tests {
  use std::vec;
@@ -932,4 +934,89 @@ mod tests {
  .await;
  }
  }
+
+ #[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
+ async fn completes_quoted_columns(pool: PgPool) {
+ let setup = r#"
+ create schema if not exists auth;
+ drop table if exists auth.quote_test_users;
+ 
+ create table auth.quote_test_users (
+ id serial primary key,
+ email text unique not null,
+ name text not null,
+ "quoted_column" text
+ );
+ "#;
+
+ pool.execute(setup).await.unwrap();
+
+ // test completion inside quoted column name
+ assert_complete_results(
+ format!(
+ r#"select "em{}" from "auth"."quote_test_users""#,
+ QueryWithCursorPosition::cursor_marker()
+ )
+ .as_str(),
+ vec![CompletionAssertion::LabelAndDesc(
+ "email".to_string(),
+ "auth.quote_test_users".to_string(),
+ )],
+ None,
+ &pool,
+ )
+ .await;
+
+ // test completion for already quoted column
+ assert_complete_results(
+ format!(
+ r#"select "quoted_col{}" from "auth"."quote_test_users""#,
+ QueryWithCursorPosition::cursor_marker()
+ )
+ .as_str(),
+ vec![CompletionAssertion::LabelAndDesc(
+ "quoted_column".to_string(),
+ "auth.quote_test_users".to_string(),
+ )],
+ None,
+ &pool,
+ )
+ .await;
+
+ // test completion with empty quotes
+ assert_complete_results(
+ format!(
+ r#"select "{}" from "auth"."quote_test_users""#,
+ QueryWithCursorPosition::cursor_marker()
+ )
+ .as_str(),
+ vec![
+ CompletionAssertion::Label("email".to_string()),
+ CompletionAssertion::Label("id".to_string()),
+ CompletionAssertion::Label("name".to_string()),
+ CompletionAssertion::Label("quoted_column".to_string()),
+ ],
+ None,
+ &pool,
+ )
+ .await;
+
+ // test completion with partially opened quote
+ assert_complete_results(
+ format!(
+ r#"select "{} from "auth"."quote_test_users""#,
+ QueryWithCursorPosition::cursor_marker()
+ )
+ .as_str(),
+ vec![
+ CompletionAssertion::Label("email".to_string()),
+ CompletionAssertion::Label("id".to_string()),
+ CompletionAssertion::Label("name".to_string()),
+ CompletionAssertion::Label("quoted_column".to_string()),
+ ],
+ None,
+ &pool,
+ )
+ .await;
+ }
 }

crates/pgt_completions/src/providers/functions.rs

@@ -4,11 +4,11 @@ use pgt_treesitter::TreesitterContext;
 use crate::{
  CompletionItemKind, CompletionText,
  builder::{CompletionBuilder, PossibleCompletionItem},
- providers::helper::get_range_to_replace,
+ providers::helper::{get_range_to_replace, with_closed_quote},
  relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-use super::helper::get_completion_text_with_schema_or_alias;
+use super::helper::with_schema_or_alias;
 
 pub fn complete_functions<'a>(
  ctx: &'a TreesitterContext,
@@ -36,9 +36,8 @@ pub fn complete_functions<'a>(
 
 fn get_completion_text(ctx: &TreesitterContext, 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());
+ let closed_quote = with_closed_quote(ctx, &func.name);
+ let mut text = with_schema_or_alias(ctx, closed_quote.as_str(), Some(func.schema.as_str()));
 
  if ctx.is_invocation {
  CompletionText {

crates/pgt_completions/src/providers/helper.rs

@@ -1,7 +1,7 @@
 use pgt_text_size::{TextRange, TextSize};
 use pgt_treesitter::TreesitterContext;
 
-use crate::{CompletionText, remove_sanitized_token};
+use crate::{is_sanitized_token_with_quote, remove_sanitized_token};
 
 pub(crate) fn find_matching_alias_for_table(
  ctx: &TreesitterContext,
@@ -19,33 +19,45 @@ pub(crate) fn get_range_to_replace(ctx: &TreesitterContext) -> TextRange {
  match ctx.node_under_cursor.as_ref() {
  Some(node) => {
  let content = ctx.get_node_under_cursor_content().unwrap_or("".into());
- let length = remove_sanitized_token(content.as_str()).len();
+ let content = content.as_str();
+
+ let length = remove_sanitized_token(content).len();
 
  let start = node.start_byte();
- let end = start + length;
+ let mut end = start + length;
+
+ if is_sanitized_token_with_quote(content) {
+ end += 1;
+ }
 
  TextRange::new(start.try_into().unwrap(), end.try_into().unwrap())
  }
  None => TextRange::empty(TextSize::new(0)),
  }
 }
 
-pub(crate) fn get_completion_text_with_schema_or_alias(
+pub(crate) fn with_schema_or_alias(
  ctx: &TreesitterContext,
  item_name: &str,
- schema_or_alias_name: &str,
-) -> Option<CompletionText> {
+ schema_or_alias_name: Option<&str>,
+) -> String {
  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
+ if schema_or_alias_name.is_none_or(|s| s == "public") || is_already_prefixed_with_schema_name {
+ item_name.to_string()
  } else {
- let range = get_range_to_replace(ctx);
+ format!("{}.{}", schema_or_alias_name.unwrap(), item_name).to_string()
+ }
+}
 
- Some(CompletionText {
- text: format!("{}.{}", schema_or_alias_name, item_name),
- range,
- is_snippet: false,
- })
+pub(crate) fn with_closed_quote(ctx: &TreesitterContext, item_name: &str) -> String {
+ let mut with_closed = String::from(item_name);
+
+ if let Some(content) = ctx.get_node_under_cursor_content() {
+ if is_sanitized_token_with_quote(content.as_str()) {
+ with_closed.push('"');
+ }
  }
+
+ with_closed
 }

crates/pgt_completions/src/providers/tables.rs

@@ -1,13 +1,15 @@
-use pgt_schema_cache::SchemaCache;
+use pgt_schema_cache::{SchemaCache, Table};
 use pgt_treesitter::TreesitterContext;
 
 use crate::{
+ CompletionText,
  builder::{CompletionBuilder, PossibleCompletionItem},
  item::CompletionItemKind,
+ providers::helper::{get_range_to_replace, with_closed_quote},
  relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
 };
 
-use super::helper::get_completion_text_with_schema_or_alias;
+use super::helper::with_schema_or_alias;
 
 pub fn complete_tables<'a>(
  ctx: &'a TreesitterContext,
@@ -34,17 +36,25 @@ pub fn complete_tables<'a>(
  description: table.schema.to_string(),
  kind: CompletionItemKind::Table,
  detail,
- completion_text: get_completion_text_with_schema_or_alias(
- ctx,
- &table.name,
- &table.schema,
- ),
+ completion_text: Some(get_completion_text(ctx, table)),
  };
 
  builder.add_item(item);
  }
 }
 
+fn get_completion_text(ctx: &TreesitterContext, table: &Table) -> CompletionText {
+ let range = get_range_to_replace(ctx);
+ let closed_quote = with_closed_quote(ctx, &table.name);
+ let text = with_schema_or_alias(ctx, closed_quote.as_str(), Some(table.schema.as_str()));
+
+ CompletionText {
+ text,
+ range,
+ is_snippet: false,
+ }
+}
+
 #[cfg(test)]
 mod tests {
 

crates/pgt_completions/src/relevance/filtering.rs

@@ -39,12 +39,35 @@ impl CompletionFilter<'_> {
  if current_node_kind.starts_with("keyword_")
  || current_node_kind == "="
  || current_node_kind == ","
- || current_node_kind == "literal"
  || current_node_kind == "ERROR"
  {
  return None;
  }
 
+ // "literal" nodes can be identfiers wrapped in quotes:
+ // `select "email" from auth.users;`
+ // Here, "email" is a literal node.
+ if current_node_kind == "literal" {
+ match self.data {
+ CompletionRelevanceData::Column(_) => match ctx.wrapping_clause_type.as_ref() {
+ Some(WrappingClause::Select)
+ | Some(WrappingClause::Where)
+ | Some(WrappingClause::Join { .. })
+ | Some(WrappingClause::Update)
+ | Some(WrappingClause::Delete)
+ | Some(WrappingClause::Insert)
+ | Some(WrappingClause::DropColumn)
+ | Some(WrappingClause::AlterColumn)
+ | Some(WrappingClause::RenameColumn)
+ | Some(WrappingClause::PolicyCheck) => {
+ // the literal is probably a column
+ }
+ _ => return None,
+ },
+ _ => return None,
+ }
+ }
+
  // No autocompletions if there are two identifiers without a separator.
  if ctx.node_under_cursor.as_ref().is_some_and(|n| match n {
  NodeUnderCursor::TsNode(node) => node.prev_sibling().is_some_and(|p| {

crates/pgt_completions/src/sanitization.rs

@@ -29,6 +29,10 @@ pub(crate) fn is_sanitized_token(txt: &str) -> bool {
  txt == SANITIZED_TOKEN || txt == SANITIZED_TOKEN_WITH_QUOTE
 }
 
+pub(crate) fn is_sanitized_token_with_quote(txt: &str) -> bool {
+ txt == SANITIZED_TOKEN_WITH_QUOTE
+}
+
 #[derive(PartialEq, Eq, Debug)]
 pub(crate) enum NodeText {
  Replaced,