initial commit

Author: juleswritescode

Cargo.lock

@@ -75,6 +75,7 @@ pg_schema_cache           = { path = "./crates/pg_schema_cache", version = "0.0.
 pg_statement_splitter     = { path = "./crates/pg_statement_splitter", version = "0.0.0" }
 pg_syntax                 = { path = "./crates/pg_syntax", version = "0.0.0" }
 pg_text_edit              = { path = "./crates/pg_text_edit", version = "0.0.0" }
+pg_treesitter_queries     = { path = "./crates/pg_treesitter_queries", version = "0.0.0" }
 pg_type_resolver          = { path = "./crates/pg_type_resolver", version = "0.0.0" }
 pg_typecheck              = { path = "./crates/pg_typecheck", version = "0.0.0" }
 pg_workspace              = { path = "./crates/pg_workspace", version = "0.0.0" }

Cargo.toml

@@ -19,6 +19,7 @@ text-size.workspace = true
 pg_schema_cache.workspace = true
 tree-sitter.workspace     = true
 tree_sitter_sql.workspace = true
+pg_treesitter_queries.workspace = true
 
 sqlx.workspace = true
 

crates/pg_completions/Cargo.toml

@@ -30,8 +30,8 @@ impl IntoIterator for CompletionResult {
     }
 }
 
-pub fn complete(params: CompletionParams) -> CompletionResult {
-    let ctx = CompletionContext::new(&params);
+pub async fn complete<'a>(params: CompletionParams<'a>) -> CompletionResult {
+    let ctx = CompletionContext::new(&params).await;
 
     let mut builder = CompletionBuilder::new();
 

crates/pg_completions/src/complete.rs

@@ -1,4 +1,7 @@
+use std::ops::Range;
+
 use pg_schema_cache::SchemaCache;
+use pg_treesitter_queries::{queries, TreeSitterQueriesExecutor};
 
 use crate::CompletionParams;
 
@@ -52,10 +55,13 @@ pub(crate) struct CompletionContext<'a> {
     pub schema_name: Option<String>,
     pub wrapping_clause_type: Option<ClauseType>,
     pub is_invocation: bool,
+    pub wrapping_statement_range: Option<Range<usize>>,
+
+    pub ts_query_executor: Option<TreeSitterQueriesExecutor<'a>>,
 }
 
 impl<'a> CompletionContext<'a> {
-    pub fn new(params: &'a CompletionParams) -> Self {
+    pub async fn new(params: &'a CompletionParams<'a>) -> Self {
         let mut ctx = Self {
             tree: params.tree,
             text: &params.text,
@@ -65,14 +71,30 @@ impl<'a> CompletionContext<'a> {
             ts_node: None,
             schema_name: None,
             wrapping_clause_type: None,
+            wrapping_statement_range: None,
             is_invocation: false,
+            ts_query_executor: None,
         };
 
         ctx.gather_tree_context();
+        ctx.dispatch_ts_queries().await;
 
         ctx
     }
 
+    async fn dispatch_ts_queries(&mut self) {
+        let tree = match self.tree.as_ref() {
+            None => return,
+            Some(t) => t,
+        };
+
+        let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), self.text);
+
+        executor.add_query_results::<queries::RelationMatch>().await;
+
+        self.ts_query_executor = Some(executor);
+    }
+
     pub fn get_ts_node_content(&self, ts_node: tree_sitter::Node<'a>) -> Option<&'a str> {
         let source = self.text;
         match ts_node.utf8_text(source.as_bytes()) {
@@ -100,36 +122,38 @@ impl<'a> CompletionContext<'a> {
          * We'll therefore adjust the cursor position such that it meets the last node of the AST.
          * `select * from use           {}` becomes `select * from use{}`.
          */
-        let current_node_kind = cursor.node().kind();
+        let current_node = cursor.node();
         while cursor.goto_first_child_for_byte(self.position).is_none() && self.position > 0 {
             self.position -= 1;
         }
 
-        self.gather_context_from_node(cursor, current_node_kind);
+        self.gather_context_from_node(cursor, current_node);
     }
 
     fn gather_context_from_node(
         &mut self,
         mut cursor: tree_sitter::TreeCursor<'a>,
-        previous_node_kind: &str,
+        previous_node: tree_sitter::Node<'a>,
     ) {
         let current_node = cursor.node();
-        let current_node_kind = current_node.kind();
 
         // prevent infinite recursion – this can happen if we only have a PROGRAM node
-        if current_node_kind == previous_node_kind {
+        if current_node.kind() == previous_node.kind() {
             self.ts_node = Some(current_node);
             return;
         }
 
-        match previous_node_kind {
-            "statement" => self.wrapping_clause_type = current_node_kind.try_into().ok(),
+        match previous_node.kind() {
+            "statement" => {
+                self.wrapping_clause_type = current_node.kind().try_into().ok();
+                self.wrapping_statement_range = Some(previous_node.byte_range());
+            }
             "invocation" => self.is_invocation = true,
 
             _ => {}
         }
 
-        match current_node_kind {
+        match current_node.kind() {
             "object_reference" => {
                 let txt = self.get_ts_node_content(current_node);
                 if let Some(txt) = txt {
@@ -159,7 +183,7 @@ impl<'a> CompletionContext<'a> {
         }
 
         cursor.goto_first_child_for_byte(self.position);
-        self.gather_context_from_node(cursor, current_node_kind);
+        self.gather_context_from_node(cursor, current_node);
     }
 }
 
@@ -179,8 +203,8 @@ mod tests {
         parser.parse(input, None).expect("Unable to parse tree")
     }
 
-    #[test]
-    fn identifies_clauses() {
+    #[tokio::test]
+    async fn identifies_clauses() {
         let test_cases = vec![
             (format!("Select {}* from users;", CURSOR_POS), "select"),
             (format!("Select * from u{};", CURSOR_POS), "from"),
@@ -220,14 +244,14 @@ mod tests {
                 schema: &pg_schema_cache::SchemaCache::new(),
             };
 
-            let ctx = CompletionContext::new(&params);
+            let ctx = CompletionContext::new(&params).await;
 
             assert_eq!(ctx.wrapping_clause_type, expected_clause.try_into().ok());
         }
     }
 
-    #[test]
-    fn identifies_schema() {
+    #[tokio::test]
+    async fn identifies_schema() {
         let test_cases = vec![
             (
                 format!("Select * from private.u{}", CURSOR_POS),
@@ -252,14 +276,14 @@ mod tests {
                 schema: &pg_schema_cache::SchemaCache::new(),
             };
 
-            let ctx = CompletionContext::new(&params);
+            let ctx = CompletionContext::new(&params).await;
 
             assert_eq!(ctx.schema_name, expected_schema.map(|f| f.to_string()));
         }
     }
 
-    #[test]
-    fn identifies_invocation() {
+    #[tokio::test]
+    async fn identifies_invocation() {
         let test_cases = vec![
             (format!("Select * from u{}sers", CURSOR_POS), false),
             (format!("Select * from u{}sers()", CURSOR_POS), true),
@@ -286,14 +310,14 @@ mod tests {
                 schema: &pg_schema_cache::SchemaCache::new(),
             };
 
-            let ctx = CompletionContext::new(&params);
+            let ctx = CompletionContext::new(&params).await;
 
             assert_eq!(ctx.is_invocation, is_invocation);
         }
     }
 
-    #[test]
-    fn does_not_fail_on_leading_whitespace() {
+    #[tokio::test]
+    async fn does_not_fail_on_leading_whitespace() {
         let cases = vec![
             format!("{}      select * from", CURSOR_POS),
             format!(" {}      select * from", CURSOR_POS),
@@ -311,7 +335,7 @@ mod tests {
                 schema: &pg_schema_cache::SchemaCache::new(),
             };
 
-            let ctx = CompletionContext::new(&params);
+            let ctx = CompletionContext::new(&params).await;
 
             let node = ctx.ts_node.map(|n| n.clone()).unwrap();
 
@@ -324,8 +348,8 @@ mod tests {
         }
     }
 
-    #[test]
-    fn does_not_fail_on_trailing_whitespace() {
+    #[tokio::test]
+    async fn does_not_fail_on_trailing_whitespace() {
         let query = format!("select * from   {}", CURSOR_POS);
 
         let (position, text) = get_text_and_position(query.as_str());
@@ -339,7 +363,7 @@ mod tests {
             schema: &pg_schema_cache::SchemaCache::new(),
         };
 
-        let ctx = CompletionContext::new(&params);
+        let ctx = CompletionContext::new(&params).await;
 
         let node = ctx.ts_node.map(|n| n.clone()).unwrap();
 
@@ -350,8 +374,8 @@ mod tests {
         );
     }
 
-    #[test]
-    fn does_not_fail_with_empty_statements() {
+    #[tokio::test]
+    async fn does_not_fail_with_empty_statements() {
         let query = format!("{}", CURSOR_POS);
 
         let (position, text) = get_text_and_position(query.as_str());
@@ -365,16 +389,16 @@ mod tests {
             schema: &pg_schema_cache::SchemaCache::new(),
         };
 
-        let ctx = CompletionContext::new(&params);
+        let ctx = CompletionContext::new(&params).await;
 
         let node = ctx.ts_node.map(|n| n.clone()).unwrap();
 
         assert_eq!(ctx.get_ts_node_content(node), Some(""));
         assert_eq!(ctx.wrapping_clause_type, None);
     }
 
-    #[test]
-    fn does_not_fail_on_incomplete_keywords() {
+    #[tokio::test]
+    async fn does_not_fail_on_incomplete_keywords() {
         //  Instead of autocompleting "FROM", we'll assume that the user
         // is selecting a certain column name, such as `frozen_account`.
         let query = format!("select * fro{}", CURSOR_POS);
@@ -390,7 +414,7 @@ mod tests {
             schema: &pg_schema_cache::SchemaCache::new(),
         };
 
-        let ctx = CompletionContext::new(&params);
+        let ctx = CompletionContext::new(&params).await;
 
         let node = ctx.ts_node.map(|n| n.clone()).unwrap();
 

crates/pg_completions/src/context.rs

@@ -45,7 +45,7 @@ mod tests {
 
         let (tree, cache) = get_test_deps(setup, &query).await;
         let params = get_test_params(&tree, &cache, &query);
-        let results = complete(params);
+        let results = complete(params).await;
 
         let CompletionItem { label, .. } = results
             .into_iter()
@@ -78,7 +78,7 @@ mod tests {
 
         let (tree, cache) = get_test_deps(setup, &query).await;
         let params = get_test_params(&tree, &cache, &query);
-        let results = complete(params);
+        let results = complete(params).await;
 
         let CompletionItem { label, kind, .. } = results
             .into_iter()
@@ -112,7 +112,7 @@ mod tests {
 
         let (tree, cache) = get_test_deps(setup, &query).await;
         let params = get_test_params(&tree, &cache, &query);
-        let results = complete(params);
+        let results = complete(params).await;
 
         let CompletionItem { label, kind, .. } = results
             .into_iter()
@@ -146,7 +146,7 @@ mod tests {
 
         let (tree, cache) = get_test_deps(setup, &query).await;
         let params = get_test_params(&tree, &cache, &query);
-        let results = complete(params);
+        let results = complete(params).await;
 
         let CompletionItem { label, kind, .. } = results
             .into_iter()

crates/pg_completions/src/providers/functions.rs

@@ -43,7 +43,7 @@ mod tests {
 
         let (tree, cache) = get_test_deps(setup, &query).await;
         let params = get_test_params(&tree, &cache, &query);
-        let results = complete(params);
+        let results = complete(params).await;
 
         assert!(!results.items.is_empty());
 
@@ -81,7 +81,7 @@ mod tests {
         for (query, expected_label) in test_cases {
             let (tree, cache) = get_test_deps(setup, &query).await;
             let params = get_test_params(&tree, &cache, &query);
-            let results = complete(params);
+            let results = complete(params).await;
 
             assert!(!results.items.is_empty());
 
@@ -126,7 +126,7 @@ mod tests {
         for (query, expected_label) in test_cases {
             let (tree, cache) = get_test_deps(setup, &query).await;
             let params = get_test_params(&tree, &cache, &query);
-            let results = complete(params);
+            let results = complete(params).await;
 
             assert!(!results.items.is_empty());
 
@@ -163,7 +163,7 @@ mod tests {
 
         let (tree, cache) = get_test_deps(setup, &query).await;
         let params = get_test_params(&tree, &cache, &query);
-        let results = complete(params);
+        let results = complete(params).await;
 
         let CompletionItem { label, kind, .. } = results
             .into_iter()

crates/pg_completions/src/providers/tables.rs

@@ -15,6 +15,10 @@ version              = "0.0.0"
 name = "tree_print"
 path = "src/bin/tree_print.rs"
 
+[[bin]]
+name = "query_debug"
+path = "src/bin/tree_query_debug.rs"
+
 [dependencies]
 anyhow = "1.0.81"
 clap   = { version = "4.5.23", features = ["derive"] }