Merge pull request #155 from juleswritescode/feat/completions-fn

feat(completions): autocomplete for function types

Author: psteinroe

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,192 +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
- .expect("Couldn't load Schema Cache");
-
- 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
- .expect("Couldn't load Schema Cache");
-
- 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
- .expect("Couldn't load SchemaCache");
-
- 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,46 @@ use pg_schema_cache::SchemaCache;
 
 use crate::CompletionParams;
 
+#[derive(Debug, PartialEq, Eq)]
+pub enum ClauseType {
+ Select,
+ Where,
+ From,
+ Update,
+ Delete,
+}
+
+impl TryFrom<&str> for ClauseType {
+ type Error = String;
+
+ fn try_from(value: &str) -> Result<Self, Self::Error> {
+ match value {
+ "select" => Ok(Self::Select),
+ "where" => Ok(Self::Where),
+ "from" => Ok(Self::From),
+ "update" => Ok(Self::Update),
+ "delete" => Ok(Self::Delete),
+ _ => {
+ let message = format!("Unimplemented ClauseType: {}", value);
+
+ // Err on tests, so we notice that we're lacking an implementation immediately.
+ if cfg!(test) {
+ panic!("{}", message);
+ }
+
+ return Err(message);
+ }
+ }
+ }
+}
+
+impl TryFrom<String> for ClauseType {
+ type Error = String;
+ fn try_from(value: String) -> Result<ClauseType, Self::Error> {
+ ClauseType::try_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 +50,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 {
+ let mut ctx = Self {
  tree: params.tree,
- text: params.text,
+ text: &params.text,
  schema_cache: params.schema,
  position: usize::from(params.position),
 
@@ -28,9 +68,9 @@ impl<'a> CompletionContext<'a> {
  is_invocation: false,
  };
 
- tree.gather_tree_context();
+ ctx.gather_tree_context();
 
- tree
+ ctx
  }
 
  pub fn get_ts_node_content(&self, ts_node: tree_sitter::Node<'a>) -> Option<&'a str> {
@@ -65,7 +105,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 = current_node_kind.try_into().ok(),
  "invocation" => self.is_invocation = true,
 
  _ => {}
@@ -84,7 +124,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 = "where".try_into().ok();
  }
 
  _ => {}
@@ -102,7 +142,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 +153,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 +189,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, expected_clause.try_into().ok());
  }
  }
 
@@ -184,7 +222,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 +257,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::*;