Skip to content

Commit ca280ab

Browse files
psteinroepsteinroe
committed
merge main
2 parents 324fc87 + 2ada420 commit ca280ab

File tree

17 files changed

+698
-90
lines changed

17 files changed

+698
-90
lines changed

‎Cargo.lock

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎biome.jsonc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
2-
"$schema": "https://biomejs.dev/schemas/1.9.4/schema.json",
2+
"$schema": "./node_modules/@biomejs/biome/configuration_schema.json",
33
"vcs": {
44
"enabled": false,
55
"clientKind": "git",
@@ -8,7 +8,7 @@
88
"files": {
99
"ignoreUnknown": false,
1010
"ignore": [],
11-
"include": ["packages/**/*"]
11+
"include": ["/packages/**/*"]
1212
},
1313
"formatter": {
1414
"enabled": true,

‎crates/pgt_completions/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ async-std = "1.12.0"
1717
pgt_text_size.workspace = true
1818

1919

20+
fuzzy-matcher = "0.3.7"
2021
pgt_schema_cache.workspace = true
2122
pgt_treesitter_queries.workspace = true
2223
schemars = { workspace = true, optional = true }

‎crates/pgt_completions/src/context.rs

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use pgt_treesitter_queries::{
88

99
use crate::sanitization::SanitizedCompletionParams;
1010

11-
#[derive(Debug, PartialEq, Eq)]
11+
#[derive(Debug, PartialEq, Eq, Hash)]
1212
pub enum WrappingClause<'a> {
1313
Select,
1414
Where,
@@ -26,6 +26,12 @@ pub(crate) enum NodeText<'a> {
2626
Original(&'a str),
2727
}
2828

29+
#[derive(PartialEq, Eq, Hash, Debug)]
30+
pub(crate) struct MentionedColumn {
31+
pub(crate) column: String,
32+
pub(crate) alias: Option<String>,
33+
}
34+
2935
/// We can map a few nodes, such as the "update" node, to actual SQL clauses.
3036
/// That gives us a lot of insight for completions.
3137
/// Other nodes, such as the "relation" node, gives us less but still
@@ -108,8 +114,8 @@ pub(crate) struct CompletionContext<'a> {
108114
pub is_in_error_node: bool,
109115

110116
pub mentioned_relations: HashMap<Option<String>, HashSet<String>>,
111-
112117
pub mentioned_table_aliases: HashMap<String, String>,
118+
pub mentioned_columns: HashMap<Option<WrappingClause<'a>>, HashSet<MentionedColumn>>,
113119
}
114120

115121
impl<'a> CompletionContext<'a> {
@@ -127,6 +133,7 @@ impl<'a> CompletionContext<'a> {
127133
is_invocation: false,
128134
mentioned_relations: HashMap::new(),
129135
mentioned_table_aliases: HashMap::new(),
136+
mentioned_columns: HashMap::new(),
130137
is_in_error_node: false,
131138
};
132139

@@ -144,36 +151,47 @@ impl<'a> CompletionContext<'a> {
144151

145152
executor.add_query_results::<queries::RelationMatch>();
146153
executor.add_query_results::<queries::TableAliasMatch>();
154+
executor.add_query_results::<queries::SelectColumnMatch>();
147155

148156
for relation_match in executor.get_iter(stmt_range) {
149157
match relation_match {
150158
QueryResult::Relation(r) => {
151159
let schema_name = r.get_schema(sql);
152160
let table_name = r.get_table(sql);
153161

154-
let current = self.mentioned_relations.get_mut(&schema_name);
155-
156-
match current {
157-
Some(c) => {
158-
c.insert(table_name);
159-
}
160-
None => {
161-
let mut new = HashSet::new();
162-
new.insert(table_name);
163-
self.mentioned_relations.insert(schema_name, new);
164-
}
165-
};
162+
if let Some(c) = self.mentioned_relations.get_mut(&schema_name) {
163+
c.insert(table_name);
164+
} else {
165+
let mut new = HashSet::new();
166+
new.insert(table_name);
167+
self.mentioned_relations.insert(schema_name, new);
168+
}
166169
}
167-
168170
QueryResult::TableAliases(table_alias_match) => {
169171
self.mentioned_table_aliases.insert(
170172
table_alias_match.get_alias(sql),
171173
table_alias_match.get_table(sql),
172174
);
173175
}
176+
QueryResult::SelectClauseColumns(c) => {
177+
let mentioned = MentionedColumn {
178+
column: c.get_column(sql),
179+
alias: c.get_alias(sql),
180+
};
174181

175-
_ => {}
176-
}
182+
if let Some(cols) = self
183+
.mentioned_columns
184+
.get_mut(&Some(WrappingClause::Select))
185+
{
186+
cols.insert(mentioned);
187+
} else {
188+
let mut new = HashSet::new();
189+
new.insert(mentioned);
190+
self.mentioned_columns
191+
.insert(Some(WrappingClause::Select), new);
192+
}
193+
}
194+
};
177195
}
178196
}
179197

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

Lines changed: 134 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,19 @@
11
use crate::{
22
CompletionItemKind,
33
builder::{CompletionBuilder, PossibleCompletionItem},
4-
context::CompletionContext,
4+
context::{CompletionContext, WrappingClause},
55
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
66
};
77

8+
use super::helper::{find_matching_alias_for_table, get_completion_text_with_schema_or_alias};
9+
810
pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut CompletionBuilder<'a>) {
911
let available_columns = &ctx.schema_cache.columns;
1012

1113
for col in available_columns {
1214
let relevance = CompletionRelevanceData::Column(col);
1315

14-
let item = PossibleCompletionItem {
16+
let mut item = PossibleCompletionItem {
1517
label: col.name.clone(),
1618
score: CompletionScore::from(relevance.clone()),
1719
filter: CompletionFilter::from(relevance),
@@ -20,6 +22,14 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio
2022
completion_text: None,
2123
};
2224

25+
// autocomplete with the alias in a join clause if we find one
26+
if matches!(ctx.wrapping_clause_type, Some(WrappingClause::Join { .. })) {
27+
item.completion_text = find_matching_alias_for_table(ctx, col.table_name.as_str())
28+
.and_then(|alias| {
29+
get_completion_text_with_schema_or_alias(ctx, col.name.as_str(), alias.as_str())
30+
});
31+
}
32+
2333
builder.add_item(item);
2434
}
2535
}
@@ -273,60 +283,50 @@ mod tests {
273283
id1 serial primary key,
274284
name1 text,
275285
address1 text,
276-
email1 text
286+
email1 text,
287+
user_settings jsonb
277288
);
278289
279290
create table public.users (
280291
id2 serial primary key,
281292
name2 text,
282293
address2 text,
283-
email2 text
294+
email2 text,
295+
settings jsonb
284296
);
285297
"#;
286298

287-
{
288-
let test_case = TestCase {
289-
message: "",
290-
query: format!(r#"select {} from users"#, CURSOR_POS),
291-
label: "suggests from table",
292-
description: "",
293-
};
294-
295-
let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
296-
let params = get_test_params(&tree, &cache, test_case.get_input_query());
297-
let results = complete(params);
298-
299-
assert_eq!(
300-
results
301-
.into_iter()
302-
.take(4)
303-
.map(|item| item.label)
304-
.collect::<Vec<String>>(),
305-
vec!["address2", "email2", "id2", "name2"]
306-
);
307-
}
308-
309-
{
310-
let test_case = TestCase {
311-
message: "",
312-
query: format!(r#"select {} from private.users"#, CURSOR_POS),
313-
label: "suggests from table",
314-
description: "",
315-
};
299+
assert_complete_results(
300+
format!(r#"select {} from users"#, CURSOR_POS).as_str(),
301+
vec![
302+
CompletionAssertion::Label("address2".into()),
303+
CompletionAssertion::Label("email2".into()),
304+
CompletionAssertion::Label("id2".into()),
305+
CompletionAssertion::Label("name2".into()),
306+
],
307+
setup,
308+
)
309+
.await;
316310

317-
let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
318-
let params = get_test_params(&tree, &cache, test_case.get_input_query());
319-
let results = complete(params);
311+
assert_complete_results(
312+
format!(r#"select {} from private.users"#, CURSOR_POS).as_str(),
313+
vec![
314+
CompletionAssertion::Label("address1".into()),
315+
CompletionAssertion::Label("email1".into()),
316+
CompletionAssertion::Label("id1".into()),
317+
CompletionAssertion::Label("name1".into()),
318+
],
319+
setup,
320+
)
321+
.await;
320322

321-
assert_eq!(
322-
results
323-
.into_iter()
324-
.take(4)
325-
.map(|item| item.label)
326-
.collect::<Vec<String>>(),
327-
vec!["address1", "email1", "id1", "name1"]
328-
);
329-
}
323+
// asserts fuzzy finding for "settings"
324+
assert_complete_results(
325+
format!(r#"select sett{} from private.users"#, CURSOR_POS).as_str(),
326+
vec![CompletionAssertion::Label("user_settings".into())],
327+
setup,
328+
)
329+
.await;
330330
}
331331

332332
#[tokio::test]
@@ -484,4 +484,93 @@ mod tests {
484484
)
485485
.await;
486486
}
487+
488+
#[tokio::test]
489+
async fn prefers_not_mentioned_columns() {
490+
let setup = r#"
491+
create schema auth;
492+
493+
create table public.one (
494+
id serial primary key,
495+
a text,
496+
b text,
497+
z text
498+
);
499+
500+
create table public.two (
501+
id serial primary key,
502+
c text,
503+
d text,
504+
e text
505+
);
506+
"#;
507+
508+
assert_complete_results(
509+
format!(
510+
"select {} from public.one o join public.two on o.id = t.id;",
511+
CURSOR_POS
512+
)
513+
.as_str(),
514+
vec![
515+
CompletionAssertion::Label("a".to_string()),
516+
CompletionAssertion::Label("b".to_string()),
517+
CompletionAssertion::Label("c".to_string()),
518+
CompletionAssertion::Label("d".to_string()),
519+
CompletionAssertion::Label("e".to_string()),
520+
],
521+
setup,
522+
)
523+
.await;
524+
525+
// "a" is already mentioned, so it jumps down
526+
assert_complete_results(
527+
format!(
528+
"select a, {} from public.one o join public.two on o.id = t.id;",
529+
CURSOR_POS
530+
)
531+
.as_str(),
532+
vec![
533+
CompletionAssertion::Label("b".to_string()),
534+
CompletionAssertion::Label("c".to_string()),
535+
CompletionAssertion::Label("d".to_string()),
536+
CompletionAssertion::Label("e".to_string()),
537+
CompletionAssertion::Label("id".to_string()),
538+
CompletionAssertion::Label("z".to_string()),
539+
CompletionAssertion::Label("a".to_string()),
540+
],
541+
setup,
542+
)
543+
.await;
544+
545+
// "id" of table one is mentioned, but table two isn't –
546+
// its priority stays up
547+
assert_complete_results(
548+
format!(
549+
"select o.id, a, b, c, d, e, {} from public.one o join public.two on o.id = t.id;",
550+
CURSOR_POS
551+
)
552+
.as_str(),
553+
vec![
554+
CompletionAssertion::LabelAndDesc(
555+
"id".to_string(),
556+
"Table: public.two".to_string(),
557+
),
558+
CompletionAssertion::Label("z".to_string()),
559+
],
560+
setup,
561+
)
562+
.await;
563+
564+
// "id" is ambiguous, so both "id" columns are lowered in priority
565+
assert_complete_results(
566+
format!(
567+
"select id, a, b, c, d, e, {} from public.one o join public.two on o.id = t.id;",
568+
CURSOR_POS
569+
)
570+
.as_str(),
571+
vec![CompletionAssertion::Label("z".to_string())],
572+
setup,
573+
)
574+
.await;
575+
}
487576
}

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use crate::{
55
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
66
};
77

8-
use super::helper::get_completion_text_with_schema;
8+
use super::helper::get_completion_text_with_schema_or_alias;
99

1010
pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut CompletionBuilder<'a>) {
1111
let available_functions = &ctx.schema_cache.functions;
@@ -19,7 +19,11 @@ pub fn complete_functions<'a>(ctx: &'a CompletionContext, builder: &mut Completi
1919
filter: CompletionFilter::from(relevance),
2020
description: format!("Schema: {}", func.schema),
2121
kind: CompletionItemKind::Function,
22-
completion_text: get_completion_text_with_schema(ctx, &func.name, &func.schema),
22+
completion_text: get_completion_text_with_schema_or_alias(
23+
ctx,
24+
&func.name,
25+
&func.schema,
26+
),
2327
};
2428

2529
builder.add_item(item);

0 commit comments

Comments
 (0)