ok

Author: juleswritescode

crates/pg_completions/src/builder.rs

@@ -1,30 +1,55 @@
-use crate::{item::CompletionItemWithRelevance, CompletionResult};
+use tower_lsp::lsp_types::CompletionItem;
+
+use crate::{item::CompletionItemWithScore, CompletionResult};
 
 pub(crate) struct CompletionBuilder {
-    items: Vec<CompletionItemWithRelevance>,
+    items: Vec<CompletionItemWithScore>,
 }
 
 impl CompletionBuilder {
     pub fn new() -> Self {
         CompletionBuilder { items: vec![] }
     }
 
-    pub fn add_item(&mut self, item: CompletionItemWithRelevance) {
+    pub fn add_item(&mut self, item: CompletionItemWithScore) {
         self.items.push(item)
     }
 
     pub fn finish(mut self) -> CompletionResult {
         self.items.sort_by(|a, b| {
-            b.score()
-                .cmp(&a.score())
+            b.score
+                .cmp(&a.score)
                 .then_with(|| a.label().cmp(&b.label()))
         });
 
         self.items.dedup_by(|a, b| a.label() == b.label());
         self.items.truncate(crate::LIMIT);
 
-        let Self { items, .. } = self;
+        let should_preselect_first_item = self.should_preselect_first_item();
+
+        let items: Vec<CompletionItem> = self
+            .items
+            .into_iter()
+            .enumerate()
+            .map(|(idx, mut item)| {
+                if idx == 0 {
+                    item.set_preselected(should_preselect_first_item);
+                }
+                item.into()
+            })
+            .collect();
 
         CompletionResult { items }
     }
+
+    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();
+
+        first.is_some_and(|f| match second {
+            Some(s) => (f.score - s.score) > 10,
+            None => true,
+        })
+    }
 }

crates/pg_completions/src/complete.rs

@@ -1,9 +1,7 @@
 use text_size::TextSize;
+use tower_lsp::lsp_types::CompletionItem;
 
-use crate::{
-    builder::CompletionBuilder, context::CompletionContext, item::CompletionItemWithRelevance,
-    providers,
-};
+use crate::{builder::CompletionBuilder, context::CompletionContext, providers};
 
 pub const LIMIT: usize = 50;
 
@@ -17,14 +15,15 @@ pub struct CompletionParams<'a> {
 
 #[derive(Debug, Default)]
 pub struct CompletionResult {
-    pub items: Vec<CompletionItemWithRelevance>,
+    pub items: Vec<CompletionItem>,
 }
 
 pub fn complete(params: CompletionParams) -> CompletionResult {
     let ctx = CompletionContext::new(&params);
     let mut builder = CompletionBuilder::new();
 
     if let Some(node) = ctx.ts_node {
+        println!("{}", node.kind());
         match node.kind() {
             "relation" => providers::complete_tables(&ctx, &mut builder),
             _ => {}
@@ -92,11 +91,12 @@ mod tests {
         )
     }
 
+    #[tokio::test]
     async fn autocompletes_table_with_schema() {
         let test_db = get_new_test_db().await;
 
         let setup = r#"
-            create schema public;
+            create schema open;
             create schema private;
 
             create table private.users (
@@ -105,7 +105,7 @@ mod tests {
                 password text
             );
 
-            create table public.user_requests (
+            create table open.user_requests (
                 id serial primary key,
                 request text,
                 send_at timestamp with time zone

crates/pg_completions/src/item.rs

@@ -1,26 +1,32 @@
-use tower_lsp::lsp_types;
+use tower_lsp::lsp_types::{self, CompletionItem};
 
 use crate::{data::CompletionItemData, relevance::CompletionRelevance};
 
 #[derive(Debug)]
-pub struct CompletionItemWithRelevance {
-    item: lsp_types::CompletionItem,
-    relevance: CompletionRelevance,
+pub struct CompletionItemWithScore {
+    pub item: lsp_types::CompletionItem,
+    pub score: i32,
 }
 
-impl CompletionItemWithRelevance {
+impl CompletionItemWithScore {
     pub(crate) fn new(data: CompletionItemData, relevance: CompletionRelevance) -> Self {
         Self {
             item: data.into(),
-            relevance,
+            score: relevance.score(),
         }
     }
 
-    pub(crate) fn score(&self) -> i32 {
-        self.relevance.score()
-    }
-
     pub(crate) fn label(&self) -> &str {
         &self.item.label
     }
+
+    pub(crate) fn set_preselected(&mut self, is_preselected: bool) {
+        self.item.preselect = Some(is_preselected)
+    }
+}
+
+impl Into<CompletionItem> for CompletionItemWithScore {
+    fn into(self) -> CompletionItem {
+        self.item
+    }
 }

crates/pg_completions/src/providers/tables.rs

@@ -1,25 +1,21 @@
-use pg_schema_cache::Table;
-
 use crate::{
     builder::CompletionBuilder, context::CompletionContext, data::CompletionItemData,
-    item::CompletionItemWithRelevance, relevance::CompletionRelevance,
+    item::CompletionItemWithScore, relevance::CompletionRelevance,
 };
 
 pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
     let available_tables = &ctx.schema_cache.tables;
 
-    let completion_items: Vec<CompletionItemWithRelevance> = available_tables
+    let completion_items: Vec<CompletionItemWithScore> = available_tables
         .iter()
-        .map(|table| to_completion_item(ctx, table))
+        .map(|table| {
+            let data = CompletionItemData::Table(table);
+            let relevance = CompletionRelevance::from_data_and_ctx(&data, ctx);
+            CompletionItemWithScore::new(data, relevance)
+        })
         .collect();
 
     for item in completion_items {
         builder.add_item(item);
     }
 }
-
-fn to_completion_item(ctx: &CompletionContext, table: &Table) -> CompletionItemWithRelevance {
-    let data = CompletionItemData::Table(table);
-    let relevance = CompletionRelevance::from_data_and_ctx(&data, ctx);
-    CompletionItemWithRelevance::new(data, relevance)
-}