add scaffolding for function completions

Author: juleswritescode

crates/pg_completions/src/complete.rs

@@ -11,7 +11,7 @@ pub const LIMIT: usize = 50;
 pub struct CompletionParams<'a> {
     pub position: TextSize,
     pub schema: &'a pg_schema_cache::SchemaCache,
-    pub text: &'a str,
+    pub text: String,
     pub tree: Option<&'a tree_sitter::Tree>,
 }
 
@@ -77,7 +77,7 @@ mod tests {
         let p = CompletionParams {
             position: ((input.len() - 1) as u32).into(),
             schema: &schema_cache,
-            text: input,
+            text: input.into(),
             tree: Some(&tree),
         };
 
@@ -136,7 +136,7 @@ mod tests {
             let p = CompletionParams {
                 position: ((input.len() - 1) as u32).into(),
                 schema: &schema_cache,
-                text: input,
+                text: input.into(),
                 tree: Some(&tree),
             };
 
@@ -199,7 +199,7 @@ mod tests {
             let p = CompletionParams {
                 position: ((input.len() - 1) as u32).into(),
                 schema: &schema_cache,
-                text: input,
+                text: input.into(),
                 tree: Some(&tree),
             };
 

crates/pg_completions/src/context.rs

@@ -18,7 +18,7 @@ impl<'a> CompletionContext<'a> {
     pub fn new(params: &'a CompletionParams) -> Self {
         let mut tree = Self {
             tree: params.tree,
-            text: params.text,
+            text: &params.text,
             schema_cache: params.schema,
             position: usize::from(params.position),
 
@@ -102,7 +102,7 @@ impl<'a> CompletionContext<'a> {
 
 #[cfg(test)]
 mod tests {
-    use crate::context::CompletionContext;
+    use crate::{context::CompletionContext, test_helper::CURSOR_POS};
 
     fn get_tree(input: &str) -> tree_sitter::Tree {
         let mut parser = tree_sitter::Parser::new();
@@ -113,8 +113,6 @@ mod tests {
         parser.parse(input, None).expect("Unable to parse tree")
     }
 
-    static CURSOR_POS: &str = "XXX";
-
     #[test]
     fn identifies_clauses() {
         let test_cases = vec![
@@ -151,7 +149,7 @@ mod tests {
             let tree = get_tree(text.as_str());
             let params = crate::CompletionParams {
                 position: (position as u32).into(),
-                text: text.as_str(),
+                text: text,
                 tree: Some(&tree),
                 schema: &pg_schema_cache::SchemaCache::new(),
             };
@@ -184,7 +182,7 @@ mod tests {
             let tree = get_tree(text.as_str());
             let params = crate::CompletionParams {
                 position: (position as u32).into(),
-                text: text.as_str(),
+                text: text,
                 tree: Some(&tree),
                 schema: &pg_schema_cache::SchemaCache::new(),
             };
@@ -219,7 +217,7 @@ mod tests {
             let tree = get_tree(text.as_str());
             let params = crate::CompletionParams {
                 position: (position as u32).into(),
-                text: text.as_str(),
+                text: text,
                 tree: Some(&tree),
                 schema: &pg_schema_cache::SchemaCache::new(),
             };

crates/pg_completions/src/item.rs

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

crates/pg_completions/src/lib.rs

@@ -4,6 +4,7 @@ mod context;
 mod item;
 mod providers;
 mod relevance;
+mod test_helper;
 
 pub use complete::*;
 pub use item::*;

crates/pg_completions/src/providers/functions.rs

@@ -0,0 +1,55 @@
+use pg_schema_cache::Function;
+
+use crate::{
+    builder::CompletionBuilder, context::CompletionContext, CompletionItem, CompletionItemKind,
+};
+
+pub fn complete_functions(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+    let available_functions = &ctx.schema_cache.functions;
+
+    let completion_items: Vec<CompletionItem> = available_functions
+        .iter()
+        .map(|foo| CompletionItem {
+            label: foo.name,
+            score: get_score(ctx, foo),
+            description: format!("Schema: {}", foo.schema),
+            preselected: None,
+            kind: CompletionItemKind::Function,
+        })
+        .collect();
+
+    for item in completion_items {
+        builder.add_item(item);
+    }
+}
+
+fn get_score(ctx: &CompletionContext, foo: &Function) -> i32 {
+    0
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        context::CompletionContext,
+        providers::complete_tables,
+        test_helper::{get_test_deps, get_test_params, CURSOR_POS},
+    };
+
+    #[tokio::test]
+    async fn completes_fn() {
+        let setup = r#"
+          create or replace function cool() returns trigger
+          begin;
+            ## Yeahhhh
+          end;
+        "#;
+
+        let query = format!("select coo{}", CURSOR_POS);
+
+        let (tree, cache, mut builder) = get_test_deps(setup, &query).await;
+        let params = get_test_params(&tree, &cache, &query);
+        let ctx = CompletionContext::new(&params);
+
+        complete_tables(&ctx, &mut builder);
+    }
+}

crates/pg_completions/src/providers/mod.rs

@@ -1,3 +1,5 @@
+mod functions;
 mod tables;
 
+pub use functions::*;
 pub use tables::*;

crates/pg_completions/src/test_helper.rs

@@ -0,0 +1,51 @@
+use pg_schema_cache::SchemaCache;
+use pg_test_utils::test_database::get_new_test_db;
+use sqlx::Executor;
+
+use crate::{builder::CompletionBuilder, CompletionParams};
+
+pub static CURSOR_POS: &str = "XXX";
+
+pub(crate) async fn get_test_deps(
+    setup: &str,
+    input: &str,
+) -> (
+    tree_sitter::Tree,
+    pg_schema_cache::SchemaCache,
+    CompletionBuilder,
+) {
+    let test_db = get_new_test_db().await;
+
+    test_db
+        .execute(setup)
+        .await
+        .expect("Failed to execute setup query");
+
+    let schema_cache = SchemaCache::load(&test_db).await;
+
+    let mut parser = tree_sitter::Parser::new();
+    parser
+        .set_language(tree_sitter_sql::language())
+        .expect("Error loading sql language");
+
+    let tree = parser.parse(input, None).unwrap();
+    let builder = CompletionBuilder::new();
+
+    (tree, schema_cache, builder)
+}
+
+pub(crate) fn get_test_params<'a>(
+    tree: &'a tree_sitter::Tree,
+    schema_cache: &'a pg_schema_cache::SchemaCache,
+    sql: &'a str,
+) -> CompletionParams<'a> {
+    let position = sql.find(CURSOR_POS).unwrap();
+    let text = sql.replace(CURSOR_POS, "");
+
+    CompletionParams {
+        position: (position as u32).into(),
+        schema: schema_cache,
+        tree: Some(tree),
+        text,
+    }
+}