Skip to content

Commit 5ba54d9

Browse files
feat(completions): prefer table/functions based on clause type
1 parent 0afe6ec commit 5ba54d9

File tree

1 file changed

+23
-1
lines changed

1 file changed

+23
-1
lines changed

crates/pg_completions/src/relevance.rs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::context::CompletionContext;
1+
use crate::context::{ClauseType, CompletionContext};
22

33
#[derive(Debug)]
44
pub(crate) enum CompletionRelevanceData<'a> {
@@ -33,6 +33,7 @@ impl<'a> CompletionRelevance<'a> {
3333
self.check_matches_query_input(ctx);
3434
self.check_if_catalog(ctx);
3535
self.check_is_invocation(ctx);
36+
self.check_matching_clause_type(ctx);
3637

3738
self.score
3839
}
@@ -60,6 +61,27 @@ impl<'a> CompletionRelevance<'a> {
6061
};
6162
}
6263

64+
fn check_matching_clause_type(&mut self, ctx: &CompletionContext) {
65+
let clause_type = match ctx.wrapping_clause_type.as_ref() {
66+
None => return,
67+
Some(ct) => ct,
68+
};
69+
70+
self.score += match self.data {
71+
CompletionRelevanceData::Table(_) => match clause_type {
72+
ClauseType::From => 5,
73+
ClauseType::Update => 15,
74+
ClauseType::Delete => 15,
75+
_ => -50,
76+
},
77+
CompletionRelevanceData::Function(_) => match clause_type {
78+
ClauseType::Select => 5,
79+
ClauseType::From => 0,
80+
_ => -50,
81+
},
82+
}
83+
}
84+
6385
fn check_is_invocation(&mut self, ctx: &CompletionContext) {
6486
self.score += match self.data {
6587
CompletionRelevanceData::Function(_) => {

0 commit comments

Comments
 (0)