feat: autocomplete schemas

Author: juleswritescode

crates/pgt_completions/benches/sanitization.rs

@@ -27,7 +27,7 @@ fn to_params<'a>(
         position: TextSize::new(pos),
         schema: &cache,
         text,
-        tree: Some(tree),
+        tree: tree,
     }
 }
 

crates/pgt_completions/src/complete.rs

@@ -4,7 +4,7 @@ use crate::{
     builder::CompletionBuilder,
     context::CompletionContext,
     item::CompletionItem,
-    providers::{complete_columns, complete_functions, complete_tables},
+    providers::{complete_columns, complete_functions, complete_schemas, complete_tables},
     sanitization::SanitizedCompletionParams,
 };
 
@@ -32,6 +32,7 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
     complete_tables(&ctx, &mut builder);
     complete_functions(&ctx, &mut builder);
     complete_columns(&ctx, &mut builder);
+    complete_schemas(&ctx, &mut builder);
 
     builder.finish()
 }

crates/pgt_completions/src/item.rs

@@ -7,6 +7,7 @@ pub enum CompletionItemKind {
     Table,
     Function,
     Column,
+    Schema,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

crates/pgt_completions/src/providers/mod.rs

@@ -1,7 +1,9 @@
 mod columns;
 mod functions;
+mod schemas;
 mod tables;
 
 pub use columns::*;
 pub use functions::*;
+pub use schemas::*;
 pub use tables::*;

crates/pgt_completions/src/providers/schemas.rs

@@ -0,0 +1,69 @@
+use crate::{
+    CompletionItem, builder::CompletionBuilder, context::CompletionContext,
+    relevance::CompletionRelevanceData,
+};
+
+pub fn complete_schemas(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+    let available_schemas = &ctx.schema_cache.schemas;
+
+    for schema in available_schemas {
+        let relevance = CompletionRelevanceData::Schema(&schema);
+
+        let item = CompletionItem {
+            label: schema.name.clone(),
+            description: "Schema".into(),
+            preselected: false,
+            kind: crate::CompletionItemKind::Schema,
+            score: relevance.get_score(ctx),
+        };
+
+        builder.add_item(item);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use crate::{
+        CompletionItem, CompletionItemKind, complete,
+        test_helper::{CURSOR_POS, get_test_deps, get_test_params},
+    };
+
+    #[tokio::test]
+    async fn autocompletes_schemas() {
+        let setup = r#"
+            create schema private;
+            create schema auth;
+            create schema internal;
+        "#;
+
+        let query = format!("select * from {}", CURSOR_POS);
+
+        let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
+        let params = get_test_params(&tree, &cache, query.as_str().into());
+        let items = complete(params);
+
+        assert!(!items.is_empty());
+
+        for item in items.iter().take(10) {
+            println!(
+                r#""{}", score: {}, kind: {:?}"#,
+                item.label, item.score, item.kind
+            );
+        }
+
+        assert_eq!(
+            items
+                .into_iter()
+                .take(4)
+                .map(|i| i.label)
+                .collect::<Vec<String>>(),
+            vec![
+                "public".to_string(), // public always preferred
+                "auth".to_string(),
+                "internal".to_string(),
+                "private".to_string()
+            ]
+        );
+    }
+}

crates/pgt_completions/src/relevance.rs

@@ -5,6 +5,7 @@ pub(crate) enum CompletionRelevanceData<'a> {
     Table(&'a pgt_schema_cache::Table),
     Function(&'a pgt_schema_cache::Function),
     Column(&'a pgt_schema_cache::Column),
+    Schema(&'a pgt_schema_cache::Schema),
 }
 
 impl CompletionRelevanceData<'_> {
@@ -58,6 +59,7 @@ impl CompletionRelevance<'_> {
             CompletionRelevanceData::Function(f) => f.name.as_str(),
             CompletionRelevanceData::Table(t) => t.name.as_str(),
             CompletionRelevanceData::Column(c) => c.name.as_str(),
+            CompletionRelevanceData::Schema(s) => s.name.as_str(),
         };
 
         if name.starts_with(content) {
@@ -97,6 +99,10 @@ impl CompletionRelevance<'_> {
                 ClauseType::Where => 10,
                 _ => -15,
             },
+            CompletionRelevanceData::Schema(_) => match clause_type {
+                ClauseType::From => 15,
+                _ => -50,
+            },
         }
     }
 
@@ -129,6 +135,7 @@ impl CompletionRelevance<'_> {
             CompletionRelevanceData::Function(f) => f.schema.as_str(),
             CompletionRelevanceData::Table(t) => t.schema.as_str(),
             CompletionRelevanceData::Column(c) => c.schema_name.as_str(),
+            CompletionRelevanceData::Schema(s) => s.name.as_str(),
         }
     }
 
@@ -168,11 +175,7 @@ impl CompletionRelevance<'_> {
     }
 
     fn check_is_user_defined(&mut self) {
-        let schema = match self.data {
-            CompletionRelevanceData::Column(c) => &c.schema_name,
-            CompletionRelevanceData::Function(f) => &f.schema,
-            CompletionRelevanceData::Table(t) => &t.schema,
-        };
+        let schema = self.get_schema_name().to_string();
 
         let system_schemas = ["pg_catalog", "information_schema", "pg_toast"];
 

crates/pgt_completions/src/sanitization.rs

@@ -213,7 +213,7 @@ mod tests {
 
     #[test]
     fn test_cursor_after_nodes() {
-        let input = "select * from";
+        let input = "select * from ";
 
         let mut parser = tree_sitter::Parser::new();
         parser

crates/pgt_lsp/src/handlers/completions.rs

@@ -50,5 +50,6 @@ fn to_lsp_types_completion_item_kind(
         pgt_completions::CompletionItemKind::Function => lsp_types::CompletionItemKind::FUNCTION,
         pgt_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS,
         pgt_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD,
+        pgt_completions::CompletionItemKind::Schema => lsp_types::CompletionItemKind::CLASS,
     }
 }

crates/pgt_schema_cache/src/lib.rs

@@ -13,4 +13,5 @@ mod versions;
 pub use columns::*;
 pub use functions::{Behavior, Function, FunctionArg, FunctionArgs};
 pub use schema_cache::SchemaCache;
+pub use schemas::Schema;
 pub use tables::{ReplicaIdentity, Table};

crates/pgt_schema_cache/src/schemas.rs

@@ -4,9 +4,9 @@ use crate::schema_cache::SchemaCacheItem;
 
 #[derive(Debug, Clone, Default)]
 pub struct Schema {
-    id: i64,
-    name: String,
-    owner: String,
+    pub id: i64,
+    pub name: String,
+    pub owner: String,
 }
 
 impl SchemaCacheItem for Schema {