Skip to content

feats(completions): complete insert, drop/alter table, ignore many situations, improve WHERE #400

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
May 24, 2025
Merged
348 changes: 279 additions & 69 deletions crates/pgt_completions/src/context/mod.rs

Large diffs are not rendered by default.

160 changes: 157 additions & 3 deletions crates/pgt_completions/src/providers/columns.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,12 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
};

// autocomplete with the alias in a join clause if we find one
if matches!(ctx.wrapping_clause_type, Some(WrappingClause::Join { .. })) {
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())
Expand All @@ -36,11 +41,13 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio

#[cfg(test)]
mod tests {
use std::vec;

use crate::{
CompletionItem, CompletionItemKind, complete,
test_helper::{
CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results, get_test_deps,
get_test_params,
CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results,
assert_no_complete_results, get_test_deps, get_test_params,
},
};

Expand Down Expand Up @@ -573,4 +580,151 @@ mod tests {
)
.await;
}

#[tokio::test]
async fn suggests_columns_in_insert_clause() {
let setup = r#"
create table instruments (
id bigint primary key generated always as identity,
name text not null,
z text
);

create table others (
id serial primary key,
a text,
b text
);
"#;

// We should prefer the instrument columns, even though they
// are lower in the alphabet

assert_complete_results(
format!("insert into instruments ({})", CURSOR_POS).as_str(),
vec![
CompletionAssertion::Label("id".to_string()),
CompletionAssertion::Label("name".to_string()),
CompletionAssertion::Label("z".to_string()),
],
setup,
)
.await;

assert_complete_results(
format!("insert into instruments (id, {})", CURSOR_POS).as_str(),
vec![
CompletionAssertion::Label("name".to_string()),
CompletionAssertion::Label("z".to_string()),
],
setup,
)
.await;

assert_complete_results(
format!("insert into instruments (id, {}, name)", CURSOR_POS).as_str(),
vec![CompletionAssertion::Label("z".to_string())],
setup,
)
.await;

// works with completed statement
assert_complete_results(
format!(
"insert into instruments (name, {}) values ('my_bass');",
CURSOR_POS
)
.as_str(),
vec![
CompletionAssertion::Label("id".to_string()),
CompletionAssertion::Label("z".to_string()),
],
setup,
)
.await;

// no completions in the values list!
assert_no_complete_results(
format!("insert into instruments (id, name) values ({})", CURSOR_POS).as_str(),
setup,
)
.await;
}

#[tokio::test]
async fn suggests_columns_in_where_clause() {
let setup = r#"
create table instruments (
id bigint primary key generated always as identity,
name text not null,
z text,
created_at timestamp with time zone default now()
);

create table others (
a text,
b text,
c text
);
"#;

assert_complete_results(
format!("select name from instruments where {} ", CURSOR_POS).as_str(),
vec![
CompletionAssertion::Label("created_at".into()),
CompletionAssertion::Label("id".into()),
CompletionAssertion::Label("name".into()),
CompletionAssertion::Label("z".into()),
],
setup,
)
.await;

assert_complete_results(
format!(
"select name from instruments where z = 'something' and created_at > {}",
CURSOR_POS
)
.as_str(),
// simply do not complete columns + schemas; functions etc. are ok
vec![
CompletionAssertion::KindNotExists(CompletionItemKind::Column),
CompletionAssertion::KindNotExists(CompletionItemKind::Schema),
],
setup,
)
.await;

// prefers not mentioned columns
assert_complete_results(
format!(
"select name from instruments where id = 'something' and {}",
CURSOR_POS
)
.as_str(),
vec![
CompletionAssertion::Label("created_at".into()),
CompletionAssertion::Label("name".into()),
CompletionAssertion::Label("z".into()),
],
setup,
)
.await;

// // uses aliases
assert_complete_results(
format!(
"select name from instruments i join others o on i.z = o.a where i.{}",
CURSOR_POS
)
.as_str(),
vec![
CompletionAssertion::Label("created_at".into()),
CompletionAssertion::Label("id".into()),
CompletionAssertion::Label("name".into()),
],
setup,
)
.await;
}
}
119 changes: 119 additions & 0 deletions crates/pgt_completions/src/providers/tables.rs
Original file line number Diff line number Diff line change
Expand Up @@ -310,4 +310,123 @@ mod tests {
)
.await;
}

#[tokio::test]
async fn suggests_tables_in_alter_and_drop_statements() {
let setup = r#"
create schema auth;

create table auth.users (
uid serial primary key,
name text not null,
email text unique not null
);

create table auth.posts (
pid serial primary key,
user_id int not null references auth.users(uid),
title text not null,
content text,
created_at timestamp default now()
);
"#;

assert_complete_results(
format!("alter table {}", CURSOR_POS).as_str(),
vec![
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
],
setup,
)
.await;

assert_complete_results(
format!("alter table if exists {}", CURSOR_POS).as_str(),
vec![
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
],
setup,
)
.await;

assert_complete_results(
format!("drop table {}", CURSOR_POS).as_str(),
vec![
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
],
setup,
)
.await;

assert_complete_results(
format!("drop table if exists {}", CURSOR_POS).as_str(),
vec![
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
],
setup,
)
.await;
}

#[tokio::test]
async fn suggests_tables_in_insert_into() {
let setup = r#"
create schema auth;

create table auth.users (
uid serial primary key,
name text not null,
email text unique not null
);
"#;

assert_complete_results(
format!("insert into {}", CURSOR_POS).as_str(),
vec![
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
],
setup,
)
.await;

assert_complete_results(
format!("insert into auth.{}", CURSOR_POS).as_str(),
vec![CompletionAssertion::LabelAndKind(
"users".into(),
CompletionItemKind::Table,
)],
setup,
)
.await;

// works with complete statement.
assert_complete_results(
format!(
"insert into {} (name, email) values ('jules', 'a@b.com');",
CURSOR_POS
)
.as_str(),
vec![
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
],
setup,
)
.await;
}
}
Loading