feat: autocomplete schemas (#340)

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,70 @@
+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::{
+        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;
+
+            -- add a table to compete against schemas
+            create table users (
+                id serial primary key,
+                name text,
+                password text
+            );
+        "#;
+
+        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());
+
+        assert_eq!(
+            items
+                .into_iter()
+                .take(5)
+                .map(|i| (i.label, i.kind))
+                .collect::<Vec<(String, CompletionItemKind)>>(),
+            vec![
+                ("public".to_string(), CompletionItemKind::Schema),
+                ("auth".to_string(), CompletionItemKind::Schema),
+                ("internal".to_string(), CompletionItemKind::Schema),
+                ("private".to_string(), CompletionItemKind::Schema),
+                ("users".to_string(), CompletionItemKind::Table),
+            ]
+        );
+    }
+}

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 => 10,
+                _ => -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

@@ -42,12 +42,28 @@ where
         let cursor_pos: usize = params.position.into();
         let mut sql = String::new();
 
-        for (idx, c) in params.text.chars().enumerate() {
-            if idx == cursor_pos {
-                sql.push_str(SANITIZED_TOKEN);
-                sql.push(' ');
+        let mut sql_iter = params.text.chars();
+
+        for idx in 0..cursor_pos + 1 {
+            match sql_iter.next() {
+                Some(c) => {
+                    if idx == cursor_pos {
+                        sql.push_str(SANITIZED_TOKEN);
+                        sql.push(' ');
+                    }
+                    sql.push(c);
+                }
+                None => {
+                    // the cursor is outside the statement,
+                    // we want to push spaces until we arrive at the cursor position.
+                    // we'll then add the SANITIZED_TOKEN
+                    if idx == cursor_pos {
+                        sql.push_str(SANITIZED_TOKEN);
+                    } else {
+                        sql.push(' ');
+                    }
+                }
             }
-            sql.push(c);
         }
 
         let mut parser = tree_sitter::Parser::new();

crates/pgt_completions/src/test_helper.rs

@@ -18,7 +18,7 @@ impl From<&str> for InputQuery {
             .expect("Insert Cursor Position into your Query.");
 
         InputQuery {
-            sql: value.replace(CURSOR_POS, ""),
+            sql: value.replace(CURSOR_POS, "").trim().to_string(),
             position,
         }
     }

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 {