feat: clause_type enum

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;
@@ -34,6 +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()
 }

crates/pg_completions/src/context.rs

@@ -2,6 +2,30 @@ use pg_schema_cache::SchemaCache;
 
 use crate::CompletionParams;
 
+#[derive(Debug, PartialEq, Eq)]
+pub enum ClauseType {
+ Select,
+ Where,
+ From,
+}
+
+impl From<&str> for ClauseType {
+ fn from(value: &str) -> Self {
+ match value {
+ "select" => Self::Select,
+ "where" => Self::Where,
+ "from" => Self::From,
+ _ => 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,7 +34,7 @@ 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,
 }
 
@@ -65,7 +89,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 +108,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());
  }
 
  _ => {}
@@ -156,7 +180,7 @@ mod tests {
 
  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()));
  }
  }
 

crates/pg_completions/src/item.rs

@@ -1,4 +1,4 @@
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum CompletionItemKind {
  Table,
  Function,

crates/pg_completions/src/providers/functions.rs

@@ -1,5 +1,3 @@
-use pg_schema_cache::Function;
-
 use crate::{
  builder::CompletionBuilder, context::CompletionContext, relevance::CompletionRelevanceData,
  CompletionItem, CompletionItemKind,
@@ -27,10 +25,9 @@ pub fn complete_functions(ctx: &CompletionContext, builder: &mut CompletionBuild
 #[cfg(test)]
 mod tests {
  use crate::{
- context::CompletionContext,
- providers::complete_functions,
+ complete,
  test_helper::{get_test_deps, get_test_params, CURSOR_POS},
- CompletionItem,
+ CompletionItem, CompletionItemKind,
  };
 
  #[tokio::test]
@@ -49,19 +46,83 @@ mod tests {
 
  let query = format!("select coo{}", CURSOR_POS);
 
- let (tree, cache, mut builder) = get_test_deps(setup, &query).await;
+ let (tree, cache) = get_test_deps(setup, &query).await;
  let params = get_test_params(&tree, &cache, &query);
- let ctx = CompletionContext::new(&params);
+ let results = complete(params);
 
- complete_functions(&ctx, &mut builder);
+ let CompletionItem { label, .. } = results
+ .into_iter()
+ .next()
+ .expect("Should return at least one completion item");
 
- let results = builder.finish();
+ assert_eq!(label, "cool");
+ }
 
- let CompletionItem { label, .. } = results
+ #[tokio::test]
+ async fn prefers_fn_if_invocation() {
+ let setup = r#"
+ create table coos (
+ id serial primary key,
+ name text
+ );
+
+ create or replace function cool() 
+ returns trigger
+ language plpgsql
+ security invoker
+ as $$
+ begin
+ raise exception 'dont matter';
+ end;
+ $$;
+ "#;
+
+ let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
+
+ let (tree, cache) = get_test_deps(setup, &query).await;
+ let params = get_test_params(&tree, &cache, &query);
+ let results = complete(params);
+
+ let CompletionItem { label, kind, .. } = results
+ .into_iter()
+ .next()
+ .expect("Should return at least one completion item");
+
+ assert_eq!(label, "cool");
+ assert_eq!(kind, CompletionItemKind::Function);
+ }
+
+ #[tokio::test]
+ async fn prefers_fn_in_select_clause() {
+ let setup = r#"
+ create table coos (
+ id serial primary key,
+ name text
+ );
+
+ create or replace function cool() 
+ returns trigger
+ language plpgsql
+ security invoker
+ as $$
+ begin
+ raise exception 'dont matter';
+ end;
+ $$;
+ "#;
+
+ let query = format!(r#"select coo{}"#, CURSOR_POS);
+
+ let (tree, cache) = get_test_deps(setup, &query).await;
+ let params = get_test_params(&tree, &cache, &query);
+ let results = complete(params);
+
+ let CompletionItem { label, kind, .. } = results
  .into_iter()
  .next()
  .expect("Should return at least one completion item");
 
  assert_eq!(label, "cool");
+ assert_eq!(kind, CompletionItemKind::Function);
  }
 }

crates/pg_completions/src/providers/tables.rs

@@ -23,3 +23,46 @@ pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder)
  builder.add_item(item);
  }
 }
+
+mod tests {
+ use crate::{
+ complete,
+ test_helper::{get_test_deps, get_test_params, CURSOR_POS},
+ CompletionItem, CompletionItemKind,
+ };
+
+ #[tokio::test]
+ async fn prefers_table_in_from_clause() {
+ let setup = r#"
+ create table coos (
+ id serial primary key,
+ name text
+ );
+
+ create or replace function cool() 
+ returns trigger
+ language plpgsql
+ security invoker
+ as $$
+ begin
+ raise exception 'dont matter';
+ end;
+ $$;
+ "#;
+
+ let query = format!(r#"select * from coo{}"#, CURSOR_POS);
+
+ let (tree, cache) = get_test_deps(setup, &query).await;
+ let params = get_test_params(&tree, &cache, &query);
+
+ let results = complete(params);
+
+ let CompletionItem { label, kind, .. } = results
+ .into_iter()
+ .next()
+ .expect("Should return at least one completion item");
+
+ assert_eq!(label, "coos");
+ assert_eq!(kind, CompletionItemKind::Table);
+ }
+}

crates/pg_completions/src/test_helper.rs

@@ -2,18 +2,14 @@ use pg_schema_cache::SchemaCache;
 use pg_test_utils::test_database::get_new_test_db;
 use sqlx::Executor;
 
-use crate::{builder::CompletionBuilder, CompletionParams};
+use crate::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,
-) {
+) -> (tree_sitter::Tree, pg_schema_cache::SchemaCache) {
  let test_db = get_new_test_db().await;
 
  test_db
@@ -29,9 +25,8 @@ pub(crate) async fn get_test_deps(
  .expect("Error loading sql language");
 
  let tree = parser.parse(input, None).unwrap();
- let builder = CompletionBuilder::new();
 
- (tree, schema_cache, builder)
+ (tree, schema_cache)
 }
 
 pub(crate) fn get_test_params<'a>(