feat(completions): autocomplete for function types

crates/pg_completions/src/builder.rs

@@ -28,7 +28,7 @@ impl CompletionBuilder {
             .enumerate()
             .map(|(idx, mut item)| {
                 if idx == 0 {
-                    item.preselected = Some(should_preselect_first_item);
+                    item.preselected = should_preselect_first_item;
                 }
                 item
             })

crates/pg_completions/src/complete.rs

@@ -1,8 +1,10 @@
 use text_size::TextSize;
 
 use crate::{
-    builder::CompletionBuilder, context::CompletionContext, item::CompletionItem,
-    providers::complete_tables,
+    builder::CompletionBuilder,
+    context::CompletionContext,
+    item::CompletionItem,
+    providers::{complete_functions, complete_tables},
 };
 
 pub const LIMIT: usize = 50;
@@ -11,7 +13,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>,
 }
 
@@ -34,186 +36,7 @@ pub fn complete(params: CompletionParams) -> CompletionResult {
     let mut builder = CompletionBuilder::new();
 
     complete_tables(&ctx, &mut builder);
+    complete_functions(&ctx, &mut builder);
 
     builder.finish()
 }
-
-#[cfg(test)]
-mod tests {
-    use pg_schema_cache::SchemaCache;
-    use pg_test_utils::test_database::*;
-
-    use sqlx::Executor;
-
-    use crate::{complete, CompletionParams};
-
-    #[tokio::test]
-    async fn autocompletes_simple_table() {
-        let test_db = get_new_test_db().await;
-
-        let setup = r#"
-            create table users (
-                id serial primary key,
-                name text,
-                password text
-            );
-        "#;
-
-        test_db
-            .execute(setup)
-            .await
-            .expect("Failed to execute setup query");
-
-        let input = "select * from u";
-
-        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 schema_cache = SchemaCache::load(&test_db).await;
-
-        let p = CompletionParams {
-            position: ((input.len() - 1) as u32).into(),
-            schema: &schema_cache,
-            text: input,
-            tree: Some(&tree),
-        };
-
-        let result = complete(p);
-
-        assert!(!result.items.is_empty());
-
-        let best_match = &result.items[0];
-
-        assert_eq!(
-            best_match.label, "users",
-            "Does not return the expected table to autocomplete: {}",
-            best_match.label
-        )
-    }
-
-    #[tokio::test]
-    async fn autocompletes_table_alphanumerically() {
-        let test_db = get_new_test_db().await;
-
-        let setup = r#"
-            create table addresses (
-                id serial primary key
-            );
-
-            create table users (
-                id serial primary key
-            );
-
-            create table emails (
-                id serial primary key
-            );
-        "#;
-
-        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 test_cases = vec![
-            ("select * from us", "users"),
-            ("select * from em", "emails"),
-            ("select * from ", "addresses"),
-        ];
-
-        for (input, expected_label) in test_cases {
-            let tree = parser.parse(input, None).unwrap();
-
-            let p = CompletionParams {
-                position: ((input.len() - 1) as u32).into(),
-                schema: &schema_cache,
-                text: input,
-                tree: Some(&tree),
-            };
-
-            let result = complete(p);
-
-            assert!(!result.items.is_empty());
-
-            let best_match = &result.items[0];
-
-            assert_eq!(
-                best_match.label, expected_label,
-                "Does not return the expected table to autocomplete: {}",
-                best_match.label
-            )
-        }
-    }
-
-    #[tokio::test]
-    async fn autocompletes_table_with_schema() {
-        let test_db = get_new_test_db().await;
-
-        let setup = r#"
-            create schema customer_support;
-            create schema private;
-
-            create table private.user_z (
-                id serial primary key,
-                name text,
-                password text
-            );
-
-            create table customer_support.user_y (
-                id serial primary key,
-                request text,
-                send_at timestamp with time zone
-            );
-        "#;
-
-        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 test_cases = vec![
-            ("select * from u", "user_y"), // user_y is preferred alphanumerically
-            ("select * from private.u", "user_z"),
-            ("select * from customer_support.u", "user_y"),
-        ];
-
-        for (input, expected_label) in test_cases {
-            let tree = parser.parse(input, None).unwrap();
-
-            let p = CompletionParams {
-                position: ((input.len() - 1) as u32).into(),
-                schema: &schema_cache,
-                text: input,
-                tree: Some(&tree),
-            };
-
-            let result = complete(p);
-
-            assert!(!result.items.is_empty());
-
-            let best_match = &result.items[0];
-
-            assert_eq!(
-                best_match.label, expected_label,
-                "Does not return the expected table to autocomplete: {}",
-                best_match.label
-            )
-        }
-    }
-}

crates/pg_completions/src/context.rs

@@ -2,6 +2,34 @@ use pg_schema_cache::SchemaCache;
 
 use crate::CompletionParams;
 
+#[derive(Debug, PartialEq, Eq)]
+pub enum ClauseType {
+    Select,
+    Where,
+    From,
+    Update,
+    Delete,
+}
+
+impl From<&str> for ClauseType {
+    fn from(value: &str) -> Self {
+        match value {
+            "select" => Self::Select,
+            "where" => Self::Where,
+            "from" => Self::From,
+            "update" => Self::Update,
+            "delete" => Self::Delete,
+            _ => panic!("Unimplemented ClauseType: {}", value),
+        }
+    }
+}
+
+impl From<String> for ClauseType {
+    fn from(value: String) -> Self {
+        ClauseType::from(value.as_str())
+    }
+}
+
 pub(crate) struct CompletionContext<'a> {
     pub ts_node: Option<tree_sitter::Node<'a>>,
     pub tree: Option<&'a tree_sitter::Tree>,
@@ -10,15 +38,15 @@ pub(crate) struct CompletionContext<'a> {
     pub position: usize,
 
     pub schema_name: Option<String>,
-    pub wrapping_clause_type: Option<String>,
+    pub wrapping_clause_type: Option<ClauseType>,
     pub is_invocation: bool,
 }
 
 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),
 
@@ -65,7 +93,7 @@ impl<'a> CompletionContext<'a> {
         let current_node_kind = current_node.kind();
 
         match previous_node_kind {
-            "statement" => self.wrapping_clause_type = Some(current_node_kind.to_string()),
+            "statement" => self.wrapping_clause_type = Some(current_node_kind.into()),
             "invocation" => self.is_invocation = true,
 
             _ => {}
@@ -84,7 +112,7 @@ impl<'a> CompletionContext<'a> {
 
             // in Treesitter, the Where clause is nested inside other clauses
             "where" => {
-                self.wrapping_clause_type = Some("where".to_string());
+                self.wrapping_clause_type = Some("where".into());
             }
 
             _ => {}
@@ -102,7 +130,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 +141,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,14 +177,14 @@ 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(),
             };
 
             let ctx = CompletionContext::new(&params);
 
-            assert_eq!(ctx.wrapping_clause_type, Some(expected_clause.to_string()));
+            assert_eq!(ctx.wrapping_clause_type, Some(expected_clause.into()));
         }
     }
 
@@ -184,7 +210,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 +245,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,13 +1,14 @@
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum CompletionItemKind {
     Table,
+    Function,
 }
 
 #[derive(Debug)]
 pub struct CompletionItem {
     pub label: String,
     pub(crate) score: i32,
     pub description: String,
-    pub preselected: Option<bool>,
+    pub preselected: bool,
     pub kind: CompletionItemKind,
 }

crates/pg_completions/src/lib.rs

@@ -5,5 +5,8 @@ mod item;
 mod providers;
 mod relevance;
 
+#[cfg(test)]
+mod test_helper;
+
 pub use complete::*;
 pub use item::*;