Skip to content

Commit b2f8f0c

Browse files
juleswritescodejuleswritescode
committed
feat(completions): filter invalid completion items (#361)
1 parent f837c1a commit b2f8f0c

File tree

10 files changed

+445
-239
lines changed

10 files changed

+445
-239
lines changed

‎crates/pgt_completions/src/builder.rs

+60-27
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,80 @@
1-
use crate::item::CompletionItem;
1+
use crate::{
2+
CompletionItemKind,
3+
context::CompletionContext,
4+
item::CompletionItem,
5+
relevance::{filtering::CompletionFilter, scoring::CompletionScore},
6+
};
27

3-
pub(crate) struct CompletionBuilder {
4-
items: Vec<CompletionItem>,
8+
pub(crate) struct PossibleCompletionItem<'a> {
9+
pub label: String,
10+
pub description: String,
11+
pub kind: CompletionItemKind,
12+
pub score: CompletionScore<'a>,
13+
pub filter: CompletionFilter<'a>,
514
}
615

7-
impl CompletionBuilder {
8-
pub fn new() -> Self {
9-
CompletionBuilder { items: vec![] }
16+
pub(crate) struct CompletionBuilder<'a> {
17+
items: Vec<PossibleCompletionItem<'a>>,
18+
ctx: &'a CompletionContext<'a>,
19+
}
20+
21+
impl<'a> CompletionBuilder<'a> {
22+
pub fn new(ctx: &'a CompletionContext) -> Self {
23+
CompletionBuilder { items: vec![], ctx }
1024
}
1125

12-
pub fn add_item(&mut self, item: CompletionItem) {
26+
pub fn add_item(&mut self, item: PossibleCompletionItem<'a>) {
1327
self.items.push(item);
1428
}
1529

16-
pub fn finish(mut self) -> Vec<CompletionItem> {
17-
self.items
18-
.sort_by(|a, b| b.score.cmp(&a.score).then_with(|| a.label.cmp(&b.label)));
30+
pub fn finish(self) -> Vec<CompletionItem> {
31+
let mut items: Vec<PossibleCompletionItem> = self
32+
.items
33+
.into_iter()
34+
.filter(|i| i.filter.is_relevant(self.ctx).is_some())
35+
.collect();
36+
37+
for item in items.iter_mut() {
38+
item.score.calc_score(self.ctx);
39+
}
1940

20-
self.items.dedup_by(|a, b| a.label == b.label);
21-
self.items.truncate(crate::LIMIT);
41+
items.sort_by(|a, b| {
42+
b.score
43+
.get_score()
44+
.cmp(&a.score.get_score())
45+
.then_with(|| a.label.cmp(&b.label))
46+
});
2247

23-
let should_preselect_first_item = self.should_preselect_first_item();
48+
items.dedup_by(|a, b| a.label == b.label);
49+
items.truncate(crate::LIMIT);
2450

25-
self.items
51+
let should_preselect_first_item = should_preselect_first_item(&items);
52+
53+
items
2654
.into_iter()
2755
.enumerate()
28-
.map(|(idx, mut item)| {
29-
if idx == 0 {
30-
item.preselected = should_preselect_first_item;
56+
.map(|(idx, item)| {
57+
let preselected = idx == 0 && should_preselect_first_item;
58+
59+
CompletionItem {
60+
description: item.description,
61+
kind: item.kind,
62+
label: item.label,
63+
preselected,
64+
score: item.score.get_score(),
3165
}
32-
item
3366
})
3467
.collect()
3568
}
69+
}
3670

37-
fn should_preselect_first_item(&mut self) -> bool {
38-
let mut items_iter = self.items.iter();
39-
let first = items_iter.next();
40-
let second = items_iter.next();
71+
fn should_preselect_first_item(items: &Vec<PossibleCompletionItem>) -> bool {
72+
let mut items_iter = items.iter();
73+
let first = items_iter.next();
74+
let second = items_iter.next();
4175

42-
first.is_some_and(|f| match second {
43-
Some(s) => (f.score - s.score) > 10,
44-
None => true,
45-
})
46-
}
76+
first.is_some_and(|f| match second {
77+
Some(s) => (f.score.get_score() - s.score.get_score()) > 10,
78+
None => true,
79+
})
4780
}

‎crates/pgt_completions/src/complete.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
2727

2828
let ctx = CompletionContext::new(&sanitized_params);
2929

30-
let mut builder = CompletionBuilder::new();
30+
let mut builder = CompletionBuilder::new(&ctx);
3131

3232
complete_tables(&ctx, &mut builder);
3333
complete_functions(&ctx, &mut builder);

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

+42-7
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,21 @@
11
use crate::{
2-
CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext,
3-
relevance::CompletionRelevanceData,
2+
CompletionItemKind,
3+
builder::{CompletionBuilder, PossibleCompletionItem},
4+
context::CompletionContext,
5+
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
46
};
57

6-
pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
8+
pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) {
79
let available_columns = &ctx.schema_cache.columns;
810

911
for col in available_columns {
10-
let item = CompletionItem {
12+
let relevance = CompletionRelevanceData::Column(col);
13+
14+
let item = PossibleCompletionItem {
1115
label: col.name.clone(),
12-
score: CompletionRelevanceData::Column(col).get_score(ctx),
16+
score: CompletionScore::from(relevance.clone()),
17+
filter: CompletionFilter::from(relevance),
1318
description: format!("Table: {}.{}", col.schema_name, col.table_name),
14-
preselected: false,
1519
kind: CompletionItemKind::Column,
1620
};
1721

@@ -22,7 +26,7 @@ pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder
2226
#[cfg(test)]
2327
mod tests {
2428
use crate::{
25-
CompletionItem, complete,
29+
CompletionItem, CompletionItemKind, complete,
2630
test_helper::{CURSOR_POS, InputQuery, get_test_deps, get_test_params},
2731
};
2832

@@ -225,6 +229,37 @@ mod tests {
225229
);
226230
}
227231

232+
#[tokio::test]
233+
async fn ignores_cols_in_from_clause() {
234+
let setup = r#"
235+
create schema private;
236+
237+
create table private.users (
238+
id serial primary key,
239+
name text,
240+
address text,
241+
email text
242+
);
243+
"#;
244+
245+
let test_case = TestCase {
246+
message: "suggests user created tables first",
247+
query: format!(r#"select * from private.{}"#, CURSOR_POS),
248+
label: "",
249+
description: "",
250+
};
251+
252+
let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
253+
let params = get_test_params(&tree, &cache, test_case.get_input_query());
254+
let results = complete(params);
255+
256+
assert!(
257+
!results
258+
.into_iter()
259+
.any(|item| item.kind == CompletionItemKind::Column)
260+
);
261+
}
262+
228263
#[tokio::test]
229264
async fn prefers_columns_of_mentioned_tables() {
230265
let setup = r#"

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

+10-6
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,21 @@
11
use crate::{
2-
CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext,
3-
relevance::CompletionRelevanceData,
2+
CompletionItemKind,
3+
builder::{CompletionBuilder, PossibleCompletionItem},
4+
context::CompletionContext,
5+
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
46
};
57

6-
pub fn complete_functions(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
8+
pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
79
let available_functions = &ctx.schema_cache.functions;
810

911
for func in available_functions {
10-
let item = CompletionItem {
12+
let relevance = CompletionRelevanceData::Function(func);
13+
14+
let item = PossibleCompletionItem {
1115
label: func.name.clone(),
12-
score: CompletionRelevanceData::Function(func).get_score(ctx),
16+
score: CompletionScore::from(relevance.clone()),
17+
filter: CompletionFilter::from(relevance),
1318
description: format!("Schema: {}", func.schema),
14-
preselected: false,
1519
kind: CompletionItemKind::Function,
1620
};
1721

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

+7-6
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,21 @@
11
use crate::{
2-
CompletionItem, builder::CompletionBuilder, context::CompletionContext,
3-
relevance::CompletionRelevanceData,
2+
builder::{CompletionBuilder, PossibleCompletionItem},
3+
context::CompletionContext,
4+
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
45
};
56

6-
pub fn complete_schemas(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
7+
pub fn complete_schemas<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
78
let available_schemas = &ctx.schema_cache.schemas;
89

910
for schema in available_schemas {
1011
let relevance = CompletionRelevanceData::Schema(schema);
1112

12-
let item = CompletionItem {
13+
let item = PossibleCompletionItem {
1314
label: schema.name.clone(),
1415
description: "Schema".into(),
15-
preselected: false,
1616
kind: crate::CompletionItemKind::Schema,
17-
score: relevance.get_score(ctx),
17+
score: CompletionScore::from(relevance.clone()),
18+
filter: CompletionFilter::from(relevance),
1819
};
1920

2021
builder.add_item(item);

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

+9-7
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,21 @@
11
use crate::{
2-
builder::CompletionBuilder,
2+
builder::{CompletionBuilder, PossibleCompletionItem},
33
context::CompletionContext,
4-
item::{CompletionItem, CompletionItemKind},
5-
relevance::CompletionRelevanceData,
4+
item::CompletionItemKind,
5+
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
66
};
77

8-
pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
8+
pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
99
let available_tables = &ctx.schema_cache.tables;
1010

1111
for table in available_tables {
12-
let item = CompletionItem {
12+
let relevance = CompletionRelevanceData::Table(table);
13+
14+
let item = PossibleCompletionItem {
1315
label: table.name.clone(),
14-
score: CompletionRelevanceData::Table(table).get_score(ctx),
16+
score: CompletionScore::from(relevance.clone()),
17+
filter: CompletionFilter::from(relevance),
1518
description: format!("Schema: {}", table.schema),
16-
preselected: false,
1719
kind: CompletionItemKind::Table,
1820
};
1921

0 commit comments

Comments
 (0)