Skip to content

fix(completions): sort items by relevance on client #362

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 30 commits into from
Apr 21, 2025
Merged
Show file tree
Hide file tree
Changes from 29 commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
e54fc93
ok
juleswritescode Apr 20, 2025
ecd0034
ignore all irrelevant schemas
juleswritescode Apr 20, 2025
be10534
fix bug
juleswritescode Apr 21, 2025
1f7ef84
ok
juleswritescode Apr 21, 2025
7731b68
fix sorting
juleswritescode Apr 21, 2025
79b0fb9
ok
juleswritescode Apr 21, 2025
5f97ae0
ok
juleswritescode Apr 21, 2025
0701724
fix(completions): cursor on dot, other issues
juleswritescode Apr 21, 2025
0847f45
Merge branch 'fix/workflow-file' of https://github.com/supabase-commu…
juleswritescode Apr 21, 2025
d2ce18a
relevance
juleswritescode Apr 21, 2025
fa61904
Merge branch 'main' of https://github.com/supabase-community/postgres…
juleswritescode Apr 21, 2025
d315a59
ok
juleswritescode Apr 21, 2025
29056ad
ok
juleswritescode Apr 21, 2025
6e73a3b
add folder
juleswritescode Apr 21, 2025
a53cc46
fixes
juleswritescode Apr 21, 2025
5065b8e
oops
juleswritescode Apr 21, 2025
3e9462f
ok
juleswritescode Apr 21, 2025
dc54e9e
Merge branch 'fix/workflow-file' of https://github.com/supabase-commu…
juleswritescode Apr 21, 2025
db49040
revert things
juleswritescode Apr 21, 2025
b7913b1
oK
juleswritescode Apr 21, 2025
e3e54e4
?
juleswritescode Apr 21, 2025
ac122a5
test belongs here
juleswritescode Apr 21, 2025
f9f2502
test goes in other PR
juleswritescode Apr 21, 2025
8ca854b
Merge branch 'feat/compl-filtering' of https://github.com/supabase-co…
juleswritescode Apr 21, 2025
46ad86a
ok
juleswritescode Apr 21, 2025
3120724
merged
juleswritescode Apr 21, 2025
28ca671
Merge branch 'feat/compl-filtering' of https://github.com/supabase-co…
juleswritescode Apr 21, 2025
cc51524
fixes
juleswritescode Apr 21, 2025
61d4a92
Merge branch 'feat/compl-filtering' of https://github.com/supabase-co…
juleswritescode Apr 21, 2025
1c660f1
Merge branch 'fix/workflow-file' into fix/compl-item-sorting
juleswritescode Apr 21, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 70 additions & 27 deletions crates/pgt_completions/src/builder.rs
Original file line number Diff line number Diff line change
@@ -1,47 +1,90 @@
use crate::item::CompletionItem;
use crate::{
CompletionItemKind,
context::CompletionContext,
item::CompletionItem,
relevance::{filtering::CompletionFilter, scoring::CompletionScore},
};

pub(crate) struct CompletionBuilder {
items: Vec<CompletionItem>,
pub(crate) struct PossibleCompletionItem<'a> {
pub label: String,
pub description: String,
pub kind: CompletionItemKind,
pub score: CompletionScore<'a>,
pub filter: CompletionFilter<'a>,
}

impl CompletionBuilder {
pub fn new() -> Self {
CompletionBuilder { items: vec![] }
pub(crate) struct CompletionBuilder<'a> {
items: Vec<PossibleCompletionItem<'a>>,
ctx: &'a CompletionContext<'a>,
}

impl<'a> CompletionBuilder<'a> {
pub fn new(ctx: &'a CompletionContext) -> Self {
CompletionBuilder { items: vec![], ctx }
}

pub fn add_item(&mut self, item: CompletionItem) {
pub fn add_item(&mut self, item: PossibleCompletionItem<'a>) {
self.items.push(item);
}

pub fn finish(mut self) -> Vec<CompletionItem> {
self.items
.sort_by(|a, b| b.score.cmp(&a.score).then_with(|| a.label.cmp(&b.label)));
pub fn finish(self) -> Vec<CompletionItem> {
let mut items: Vec<PossibleCompletionItem> = self
.items
.into_iter()
.filter(|i| i.filter.is_relevant(self.ctx).is_some())
.collect();

for item in items.iter_mut() {
item.score.calc_score(self.ctx);
}

items.sort_by(|a, b| {
b.score
.get_score()
.cmp(&a.score.get_score())
.then_with(|| a.label.cmp(&b.label))
});

self.items.dedup_by(|a, b| a.label == b.label);
self.items.truncate(crate::LIMIT);
items.dedup_by(|a, b| a.label == b.label);
items.truncate(crate::LIMIT);

let should_preselect_first_item = self.should_preselect_first_item();
let should_preselect_first_item = should_preselect_first_item(&items);

self.items
/*
* LSP Clients themselves sort the completion items.
* They'll use the `sort_text` property if present (or fallback to the `label`).
* Since our items are already sorted, we're 'hijacking' the sort_text.
* We're simply adding the index of the item, padded by zeroes to the max length.
*/
let max_padding = items.len().to_string().len();

items
.into_iter()
.enumerate()
.map(|(idx, mut item)| {
if idx == 0 {
item.preselected = should_preselect_first_item;
.map(|(idx, item)| {
let preselected = idx == 0 && should_preselect_first_item;

CompletionItem {
description: item.description,
kind: item.kind,
label: item.label,
preselected,

// wonderous Rust syntax ftw
sort_text: format!("{:0>padding$}", idx, padding = max_padding),
}
item
})
.collect()
}
}

fn should_preselect_first_item(&mut self) -> bool {
let mut items_iter = self.items.iter();
let first = items_iter.next();
let second = items_iter.next();
fn should_preselect_first_item(items: &Vec<PossibleCompletionItem>) -> bool {
let mut items_iter = items.iter();
let first = items_iter.next();
let second = items_iter.next();

first.is_some_and(|f| match second {
Some(s) => (f.score - s.score) > 10,
None => true,
})
}
first.is_some_and(|f| match second {
Some(s) => (f.score.get_score() - s.score.get_score()) > 10,
None => true,
})
}
2 changes: 1 addition & 1 deletion crates/pgt_completions/src/complete.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {

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

let mut builder = CompletionBuilder::new();
let mut builder = CompletionBuilder::new(&ctx);

complete_tables(&ctx, &mut builder);
complete_functions(&ctx, &mut builder);
Expand Down
18 changes: 17 additions & 1 deletion crates/pgt_completions/src/item.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::fmt::Display;

use serde::{Deserialize, Serialize};

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
Expand All @@ -10,12 +12,26 @@ pub enum CompletionItemKind {
Schema,
}

impl Display for CompletionItemKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let txt = match self {
CompletionItemKind::Table => "Table",
CompletionItemKind::Function => "Function",
CompletionItemKind::Column => "Column",
CompletionItemKind::Schema => "Schema",
};

write!(f, "{txt}")
}
}

#[derive(Debug, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct CompletionItem {
pub label: String,
pub(crate) score: i32,
pub description: String,
pub preselected: bool,
pub kind: CompletionItemKind,
/// String used for sorting by LSP clients.
pub sort_text: String,
}
114 changes: 107 additions & 7 deletions crates/pgt_completions/src/providers/columns.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,21 @@
use crate::{
CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext,
relevance::CompletionRelevanceData,
CompletionItemKind,
builder::{CompletionBuilder, PossibleCompletionItem},
context::CompletionContext,
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
};

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

for col in available_columns {
let item = CompletionItem {
let relevance = CompletionRelevanceData::Column(col);

let item = PossibleCompletionItem {
label: col.name.clone(),
score: CompletionRelevanceData::Column(col).get_score(ctx),
score: CompletionScore::from(relevance.clone()),
filter: CompletionFilter::from(relevance),
description: format!("Table: {}.{}", col.schema_name, col.table_name),
preselected: false,
kind: CompletionItemKind::Column,
};

Expand All @@ -22,7 +26,7 @@ pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder
#[cfg(test)]
mod tests {
use crate::{
CompletionItem, complete,
CompletionItem, CompletionItemKind, complete,
test_helper::{CURSOR_POS, InputQuery, get_test_deps, get_test_params},
};

Expand Down Expand Up @@ -224,4 +228,100 @@ mod tests {
"`email` not present in first four completion items."
);
}

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

create table private.users (
id serial primary key,
name text,
address text,
email text
);
"#;

let test_case = TestCase {
message: "suggests user created tables first",
query: format!(r#"select * from private.{}"#, CURSOR_POS),
label: "",
description: "",
};

let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
let params = get_test_params(&tree, &cache, test_case.get_input_query());
let results = complete(params);

assert!(
!results
.into_iter()
.any(|item| item.kind == CompletionItemKind::Column)
);
}

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

create table private.users (
id1 serial primary key,
name1 text,
address1 text,
email1 text
);

create table public.users (
id2 serial primary key,
name2 text,
address2 text,
email2 text
);
"#;

{
let test_case = TestCase {
message: "",
query: format!(r#"select {} from users"#, CURSOR_POS),
label: "suggests from table",
description: "",
};

let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
let params = get_test_params(&tree, &cache, test_case.get_input_query());
let results = complete(params);

assert_eq!(
results
.into_iter()
.take(4)
.map(|item| item.label)
.collect::<Vec<String>>(),
vec!["address2", "email2", "id2", "name2"]
);
}

{
let test_case = TestCase {
message: "",
query: format!(r#"select {} from private.users"#, CURSOR_POS),
label: "suggests from table",
description: "",
};

let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
let params = get_test_params(&tree, &cache, test_case.get_input_query());
let results = complete(params);

assert_eq!(
results
.into_iter()
.take(4)
.map(|item| item.label)
.collect::<Vec<String>>(),
vec!["address1", "email1", "id1", "name1"]
);
}
}
}
16 changes: 10 additions & 6 deletions crates/pgt_completions/src/providers/functions.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,21 @@
use crate::{
CompletionItem, CompletionItemKind, builder::CompletionBuilder, context::CompletionContext,
relevance::CompletionRelevanceData,
CompletionItemKind,
builder::{CompletionBuilder, PossibleCompletionItem},
context::CompletionContext,
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
};

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

for func in available_functions {
let item = CompletionItem {
let relevance = CompletionRelevanceData::Function(func);

let item = PossibleCompletionItem {
label: func.name.clone(),
score: CompletionRelevanceData::Function(func).get_score(ctx),
score: CompletionScore::from(relevance.clone()),
filter: CompletionFilter::from(relevance),
description: format!("Schema: {}", func.schema),
preselected: false,
kind: CompletionItemKind::Function,
};

Expand Down
13 changes: 7 additions & 6 deletions crates/pgt_completions/src/providers/schemas.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,21 @@
use crate::{
CompletionItem, builder::CompletionBuilder, context::CompletionContext,
relevance::CompletionRelevanceData,
builder::{CompletionBuilder, PossibleCompletionItem},
context::CompletionContext,
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
};

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

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

let item = CompletionItem {
let item = PossibleCompletionItem {
label: schema.name.clone(),
description: "Schema".into(),
preselected: false,
kind: crate::CompletionItemKind::Schema,
score: relevance.get_score(ctx),
score: CompletionScore::from(relevance.clone()),
filter: CompletionFilter::from(relevance),
};

builder.add_item(item);
Expand Down
16 changes: 9 additions & 7 deletions crates/pgt_completions/src/providers/tables.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,21 @@
use crate::{
builder::CompletionBuilder,
builder::{CompletionBuilder, PossibleCompletionItem},
context::CompletionContext,
item::{CompletionItem, CompletionItemKind},
relevance::CompletionRelevanceData,
item::CompletionItemKind,
relevance::{CompletionRelevanceData, filtering::CompletionFilter, scoring::CompletionScore},
};

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

for table in available_tables {
let item = CompletionItem {
let relevance = CompletionRelevanceData::Table(table);

let item = PossibleCompletionItem {
label: table.name.clone(),
score: CompletionRelevanceData::Table(table).get_score(ctx),
score: CompletionScore::from(relevance.clone()),
filter: CompletionFilter::from(relevance),
description: format!("Schema: {}", table.schema),
preselected: false,
kind: CompletionItemKind::Table,
};

Expand Down
Loading
Loading