Skip to content

Commit 62ce45d

Browse files
juleswritescodejuleswritescode
authored
feats(completions): complete insert, drop/alter table, ignore many situations, improve WHERE (#400)
1 parent 74f913e commit 62ce45d

File tree

10 files changed

+1163
-131
lines changed

10 files changed

+1163
-131
lines changed

‎crates/pgt_completions/src/context/mod.rs

Lines changed: 279 additions & 69 deletions
Large diffs are not rendered by default.

‎crates/pgt_completions/src/providers/columns.rs

Lines changed: 157 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,12 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
2424
};
2525

2626
// autocomplete with the alias in a join clause if we find one
27-
if matches!(ctx.wrapping_clause_type, Some(WrappingClause::Join { .. })) {
27+
if matches!(
28+
ctx.wrapping_clause_type,
29+
Some(WrappingClause::Join { .. })
30+
| Some(WrappingClause::Where)
31+
| Some(WrappingClause::Select)
32+
) {
2833
item.completion_text = find_matching_alias_for_table(ctx, col.table_name.as_str())
2934
.and_then(|alias| {
3035
get_completion_text_with_schema_or_alias(ctx, col.name.as_str(), alias.as_str())
@@ -37,11 +42,13 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
3742

3843
#[cfg(test)]
3944
mod tests {
45+
use std::vec;
46+
4047
use crate::{
4148
CompletionItem, CompletionItemKind, complete,
4249
test_helper::{
43-
CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results, get_test_deps,
44-
get_test_params,
50+
CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results,
51+
assert_no_complete_results, get_test_deps, get_test_params,
4552
},
4653
};
4754

@@ -574,4 +581,151 @@ mod tests {
574581
)
575582
.await;
576583
}
584+
585+
#[tokio::test]
586+
async fn suggests_columns_in_insert_clause() {
587+
let setup = r#"
588+
create table instruments (
589+
id bigint primary key generated always as identity,
590+
name text not null,
591+
z text
592+
);
593+
594+
create table others (
595+
id serial primary key,
596+
a text,
597+
b text
598+
);
599+
"#;
600+
601+
// We should prefer the instrument columns, even though they
602+
// are lower in the alphabet
603+
604+
assert_complete_results(
605+
format!("insert into instruments ({})", CURSOR_POS).as_str(),
606+
vec![
607+
CompletionAssertion::Label("id".to_string()),
608+
CompletionAssertion::Label("name".to_string()),
609+
CompletionAssertion::Label("z".to_string()),
610+
],
611+
setup,
612+
)
613+
.await;
614+
615+
assert_complete_results(
616+
format!("insert into instruments (id, {})", CURSOR_POS).as_str(),
617+
vec![
618+
CompletionAssertion::Label("name".to_string()),
619+
CompletionAssertion::Label("z".to_string()),
620+
],
621+
setup,
622+
)
623+
.await;
624+
625+
assert_complete_results(
626+
format!("insert into instruments (id, {}, name)", CURSOR_POS).as_str(),
627+
vec![CompletionAssertion::Label("z".to_string())],
628+
setup,
629+
)
630+
.await;
631+
632+
// works with completed statement
633+
assert_complete_results(
634+
format!(
635+
"insert into instruments (name, {}) values ('my_bass');",
636+
CURSOR_POS
637+
)
638+
.as_str(),
639+
vec![
640+
CompletionAssertion::Label("id".to_string()),
641+
CompletionAssertion::Label("z".to_string()),
642+
],
643+
setup,
644+
)
645+
.await;
646+
647+
// no completions in the values list!
648+
assert_no_complete_results(
649+
format!("insert into instruments (id, name) values ({})", CURSOR_POS).as_str(),
650+
setup,
651+
)
652+
.await;
653+
}
654+
655+
#[tokio::test]
656+
async fn suggests_columns_in_where_clause() {
657+
let setup = r#"
658+
create table instruments (
659+
id bigint primary key generated always as identity,
660+
name text not null,
661+
z text,
662+
created_at timestamp with time zone default now()
663+
);
664+
665+
create table others (
666+
a text,
667+
b text,
668+
c text
669+
);
670+
"#;
671+
672+
assert_complete_results(
673+
format!("select name from instruments where {} ", CURSOR_POS).as_str(),
674+
vec![
675+
CompletionAssertion::Label("created_at".into()),
676+
CompletionAssertion::Label("id".into()),
677+
CompletionAssertion::Label("name".into()),
678+
CompletionAssertion::Label("z".into()),
679+
],
680+
setup,
681+
)
682+
.await;
683+
684+
assert_complete_results(
685+
format!(
686+
"select name from instruments where z = 'something' and created_at > {}",
687+
CURSOR_POS
688+
)
689+
.as_str(),
690+
// simply do not complete columns + schemas; functions etc. are ok
691+
vec![
692+
CompletionAssertion::KindNotExists(CompletionItemKind::Column),
693+
CompletionAssertion::KindNotExists(CompletionItemKind::Schema),
694+
],
695+
setup,
696+
)
697+
.await;
698+
699+
// prefers not mentioned columns
700+
assert_complete_results(
701+
format!(
702+
"select name from instruments where id = 'something' and {}",
703+
CURSOR_POS
704+
)
705+
.as_str(),
706+
vec![
707+
CompletionAssertion::Label("created_at".into()),
708+
CompletionAssertion::Label("name".into()),
709+
CompletionAssertion::Label("z".into()),
710+
],
711+
setup,
712+
)
713+
.await;
714+
715+
// // uses aliases
716+
assert_complete_results(
717+
format!(
718+
"select name from instruments i join others o on i.z = o.a where i.{}",
719+
CURSOR_POS
720+
)
721+
.as_str(),
722+
vec![
723+
CompletionAssertion::Label("created_at".into()),
724+
CompletionAssertion::Label("id".into()),
725+
CompletionAssertion::Label("name".into()),
726+
],
727+
setup,
728+
)
729+
.await;
730+
}
577731
}

‎crates/pgt_completions/src/providers/tables.rs

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -311,4 +311,123 @@ mod tests {
311311
)
312312
.await;
313313
}
314+
315+
#[tokio::test]
316+
async fn suggests_tables_in_alter_and_drop_statements() {
317+
let setup = r#"
318+
create schema auth;
319+
320+
create table auth.users (
321+
uid serial primary key,
322+
name text not null,
323+
email text unique not null
324+
);
325+
326+
create table auth.posts (
327+
pid serial primary key,
328+
user_id int not null references auth.users(uid),
329+
title text not null,
330+
content text,
331+
created_at timestamp default now()
332+
);
333+
"#;
334+
335+
assert_complete_results(
336+
format!("alter table {}", CURSOR_POS).as_str(),
337+
vec![
338+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
339+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
340+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
341+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
342+
],
343+
setup,
344+
)
345+
.await;
346+
347+
assert_complete_results(
348+
format!("alter table if exists {}", CURSOR_POS).as_str(),
349+
vec![
350+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
351+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
352+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
353+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
354+
],
355+
setup,
356+
)
357+
.await;
358+
359+
assert_complete_results(
360+
format!("drop table {}", CURSOR_POS).as_str(),
361+
vec![
362+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
363+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
364+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
365+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
366+
],
367+
setup,
368+
)
369+
.await;
370+
371+
assert_complete_results(
372+
format!("drop table if exists {}", CURSOR_POS).as_str(),
373+
vec![
374+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
375+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
376+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join
377+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
378+
],
379+
setup,
380+
)
381+
.await;
382+
}
383+
384+
#[tokio::test]
385+
async fn suggests_tables_in_insert_into() {
386+
let setup = r#"
387+
create schema auth;
388+
389+
create table auth.users (
390+
uid serial primary key,
391+
name text not null,
392+
email text unique not null
393+
);
394+
"#;
395+
396+
assert_complete_results(
397+
format!("insert into {}", CURSOR_POS).as_str(),
398+
vec![
399+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
400+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
401+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
402+
],
403+
setup,
404+
)
405+
.await;
406+
407+
assert_complete_results(
408+
format!("insert into auth.{}", CURSOR_POS).as_str(),
409+
vec![CompletionAssertion::LabelAndKind(
410+
"users".into(),
411+
CompletionItemKind::Table,
412+
)],
413+
setup,
414+
)
415+
.await;
416+
417+
// works with complete statement.
418+
assert_complete_results(
419+
format!(
420+
"insert into {} (name, email) values ('jules', 'a@b.com');",
421+
CURSOR_POS
422+
)
423+
.as_str(),
424+
vec![
425+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
426+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
427+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
428+
],
429+
setup,
430+
)
431+
.await;
432+
}
314433
}

0 commit comments

Comments
 (0)