fix: user's cursor can be offset for autocompletion

juleswritescode · juleswritescode · commit 8521f6da0a5c · 2024-12-17T09:29:29.000+01:00
diff --git a/crates/pg_completions/src/context.rs b/crates/pg_completions/src/context.rs
@@ -18,7 +18,7 @@ impl TryFrom<&str> for ClauseType {
         match value {
             "select" => Ok(Self::Select),
             "where" => Ok(Self::Where),
-            "from" => Ok(Self::From),
+            "from" | "keyword_from" => Ok(Self::From),
             "update" => Ok(Self::Update),
             "delete" => Ok(Self::Delete),
             _ => {
@@ -88,13 +88,22 @@ impl<'a> CompletionContext<'a> {
 
         let mut cursor = self.tree.as_ref().unwrap().root_node().walk();
 
-        // go to the statement node that matches the position
+        /*
+         * The head node of any treesitter tree is always the "PROGRAM" node.
+         *
+         * We want to enter the next layer and focus on the child node that matches the user's cursor position.
+         * If there is no node under the users position, however, the cursor won't enter the next level – it
+         * will stay on the Program node.
+         *
+         * This might lead to an unexpected context or infinite recursion.
+         *
+         * 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();
-
-        dbg!(current_node_kind);
-        dbg!(self.position);
-
-        cursor.goto_first_child_for_byte(self.position);
+        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);
     }
@@ -107,9 +116,7 @@ impl<'a> CompletionContext<'a> {
         let current_node = cursor.node();
         let current_node_kind = current_node.kind();
 
-        println!("inside..");
-        dbg!(current_node_kind);
-
+        // prevent infinite recursion – this can happen if we only have a PROGRAM node
         if current_node_kind == previous_node_kind {
             self.ts_node = Some(current_node);
             return;
@@ -138,9 +145,14 @@ impl<'a> CompletionContext<'a> {
                 self.wrapping_clause_type = "where".try_into().ok();
             }
 
+            "keyword_from" => {
+                self.wrapping_clause_type = "keyword_from".try_into().ok();
+            }
+
             _ => {}
         }
 
+        // We have arrived at the leaf node
         if current_node.child_count() == 0 {
             self.ts_node = Some(current_node);
             return;
@@ -153,7 +165,10 @@ impl<'a> CompletionContext<'a> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{context::CompletionContext, test_helper::CURSOR_POS};
+    use crate::{
+        context::{ClauseType, CompletionContext},
+        test_helper::{get_text_and_position, CURSOR_POS},
+    };
 
     fn get_tree(input: &str) -> tree_sitter::Tree {
         let mut parser = tree_sitter::Parser::new();
@@ -193,11 +208,11 @@ mod tests {
             ),
         ];
 
-        for (text, expected_clause) in test_cases {
-            let position = text.find(CURSOR_POS).unwrap();
-            let text = text.replace(CURSOR_POS, "");
+        for (query, expected_clause) in test_cases {
+            let (position, text) = get_text_and_position(query.as_str());
 
             let tree = get_tree(text.as_str());
+
             let params = crate::CompletionParams {
                 position: (position as u32).into(),
                 text: text,
@@ -226,9 +241,8 @@ mod tests {
             (format!("Select * from u{}sers()", CURSOR_POS), None),
         ];
 
-        for (text, expected_schema) in test_cases {
-            let position = text.find(CURSOR_POS).unwrap();
-            let text = text.replace(CURSOR_POS, "");
+        for (query, expected_schema) in test_cases {
+            let (position, text) = get_text_and_position(query.as_str());
 
             let tree = get_tree(text.as_str());
             let params = crate::CompletionParams {
@@ -261,9 +275,8 @@ mod tests {
             ),
         ];
 
-        for (text, is_invocation) in test_cases {
-            let position = text.find(CURSOR_POS).unwrap();
-            let text = text.replace(CURSOR_POS, "");
+        for (query, is_invocation) in test_cases {
+            let (position, text) = get_text_and_position(query.as_str());
 
             let tree = get_tree(text.as_str());
             let params = crate::CompletionParams {
@@ -280,11 +293,42 @@ mod tests {
     }
 
     #[test]
-    fn get_ts_node_content_does_not_fail_on_error_nodes() {
-        let query = format!("select * from {}", CURSOR_POS);
+    fn does_not_fail_on_leading_whitespace() {
+        let cases = vec![
+            format!("{}      select * from", CURSOR_POS),
+            format!(" {}      select * from", CURSOR_POS),
+        ];
+
+        for query in cases {
+            let (position, text) = get_text_and_position(query.as_str());
 
-        let position = query.find(CURSOR_POS).unwrap();
-        let text = query.replace(CURSOR_POS, "");
+            let tree = get_tree(text.as_str());
+
+            let params = crate::CompletionParams {
+                position: (position as u32).into(),
+                text: text,
+                tree: Some(&tree),
+                schema: &pg_schema_cache::SchemaCache::new(),
+            };
+
+            let ctx = CompletionContext::new(&params);
+
+            let node = ctx.ts_node.map(|n| n.clone()).unwrap();
+
+            assert_eq!(ctx.get_ts_node_content(node), Some("select"));
+
+            assert_eq!(
+                ctx.wrapping_clause_type,
+                Some(crate::context::ClauseType::Select)
+            );
+        }
+    }
+
+    #[test]
+    fn does_not_fail_on_trailing_whitespace() {
+        let query = format!("select * from   {}", CURSOR_POS);
+
+        let (position, text) = get_text_and_position(query.as_str());
 
         let tree = get_tree(text.as_str());
 
@@ -297,10 +341,60 @@ mod tests {
 
         let ctx = CompletionContext::new(&params);
 
-        let node = ctx.ts_node.map(|n| n.clone());
+        let node = ctx.ts_node.map(|n| n.clone()).unwrap();
+
+        assert_eq!(ctx.get_ts_node_content(node), Some("from"));
+        assert_eq!(
+            ctx.wrapping_clause_type,
+            Some(crate::context::ClauseType::From)
+        );
+    }
+
+    #[test]
+    fn does_not_fail_with_empty_statements() {
+        let query = format!("{}", CURSOR_POS);
+
+        let (position, text) = get_text_and_position(query.as_str());
+
+        let tree = get_tree(text.as_str());
+
+        let params = crate::CompletionParams {
+            position: (position as u32).into(),
+            text: text,
+            tree: Some(&tree),
+            schema: &pg_schema_cache::SchemaCache::new(),
+        };
+
+        let ctx = CompletionContext::new(&params);
+
+        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() {
+        //  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);
+
+        let (position, text) = get_text_and_position(query.as_str());
+
+        let tree = get_tree(text.as_str());
+
+        let params = crate::CompletionParams {
+            position: (position as u32).into(),
+            text: text,
+            tree: Some(&tree),
+            schema: &pg_schema_cache::SchemaCache::new(),
+        };
+
+        let ctx = CompletionContext::new(&params);
 
-        println!("node kind: {}", node.as_ref().unwrap().kind());
+        let node = ctx.ts_node.map(|n| n.clone()).unwrap();
 
-        ctx.get_ts_node_content(node.unwrap());
+        assert_eq!(ctx.get_ts_node_content(node), Some("fro"));
+        assert_eq!(ctx.wrapping_clause_type, Some(ClauseType::Select));
     }
 }
diff --git a/crates/pg_completions/src/test_helper.rs b/crates/pg_completions/src/test_helper.rs
@@ -31,16 +31,25 @@ pub(crate) async fn get_test_deps(
     (tree, schema_cache)
 }
 
+pub(crate) fn get_text_and_position(sql: &str) -> (usize, String) {
+    // the cursor is to the left of the `CURSOR_POS`
+    let position = sql
+        .find(|c| c == CURSOR_POS)
+        .expect("Please insert the CURSOR_POS into your query.")
+        .checked_sub(1)
+        .unwrap_or(0);
+
+    let text = sql.replace(CURSOR_POS, "");
+
+    (position, text)
+}
+
 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(|c| c == CURSOR_POS)
-        .expect("Please insert the CURSOR_POS into your query.");
-
-    let text = sql.replace(CURSOR_POS, "");
+    let (position, text) = get_text_and_position(sql);
 
     CompletionParams {
         position: (position as u32).into(),
diff --git a/crates/pg_schema_cache/src/queries/functions.sql b/crates/pg_schema_cache/src/queries/functions.sql
@@ -1,6 +1,15 @@
 with functions as (
   select
-    *,
+    oid,
+    proname,
+    prosrc,
+    prorettype,
+    proretset,
+    provolatile,
+    prosecdef,
+    prolang,
+    pronamespace,
+    proconfig,
     -- proargmodes is null when all arg modes are IN
     coalesce(
       p.proargmodes,
@@ -29,23 +38,23 @@ with functions as (
     p.prokind = 'f'
 )
 select
-  f.oid :: int8 as id,
-  n.nspname as schema,
-  f.proname as name,
-  l.lanname as language,
+  f.oid :: int8 as "id!",
+  n.nspname as "schema!",
+  f.proname as "name!",
+  l.lanname as "language!",
   case
-    when l.lanname = 'internal' then ''
+    when l.lanname = 'internal' then null
     else f.prosrc
-  end as definition,
+  end as body,
   case
-    when l.lanname = 'internal' then f.prosrc
+    when l.lanname = 'internal' then null
     else pg_get_functiondef(f.oid)
-  end as complete_statement,
+  end as definition,
   coalesce(f_args.args, '[]') as args,
-  pg_get_function_arguments(f.oid) as argument_types,
-  pg_get_function_identity_arguments(f.oid) as identity_argument_types,
-  f.prorettype :: int8 as return_type_id,
-  pg_get_function_result(f.oid) as return_type,
+  nullif(pg_get_function_arguments(f.oid), '') as argument_types,
+  nullif(pg_get_function_identity_arguments(f.oid), '') as identity_argument_types,
+  f.prorettype :: int8 as "return_type_id!",
+  pg_get_function_result(f.oid) as "return_type!",
   nullif(rt.typrelid :: int8, 0) as return_type_relation_id,
   f.proretset as is_set_returning_function,
   case
