fix(completions): cursor on dot, other issues

crates/pgt_completions/src/providers/columns.rs

@@ -224,4 +224,69 @@ mod tests {
             "`email` not present in first four completion items."
         );
     }
+
+    #[tokio::test]
+    async fn prefers_columns_of_mentioned_tables() {
+        let setup = r#"
+        create schema private;
+
+        create table private.users (
+            id1 serial primary key,
+            name1 text,
+            address1 text,
+            email1 text
+        );
+
+        create table public.users (
+            id2 serial primary key,
+            name2 text,
+            address2 text,
+            email2 text
+        );
+    "#;
+
+        {
+            let test_case = TestCase {
+                message: "",
+                query: format!(r#"select {} from users"#, CURSOR_POS),
+                label: "suggests from table",
+                description: "",
+            };
+
+            let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
+            let params = get_test_params(&tree, &cache, test_case.get_input_query());
+            let results = complete(params);
+
+            assert_eq!(
+                results
+                    .into_iter()
+                    .take(4)
+                    .map(|item| item.label)
+                    .collect::<Vec<String>>(),
+                vec!["address2", "email2", "id2", "name2"]
+            );
+        }
+
+        {
+            let test_case = TestCase {
+                message: "",
+                query: format!(r#"select {} from private.users"#, CURSOR_POS),
+                label: "suggests from table",
+                description: "",
+            };
+
+            let (tree, cache) = get_test_deps(setup, test_case.get_input_query()).await;
+            let params = get_test_params(&tree, &cache, test_case.get_input_query());
+            let results = complete(params);
+
+            assert_eq!(
+                results
+                    .into_iter()
+                    .take(4)
+                    .map(|item| item.label)
+                    .collect::<Vec<String>>(),
+                vec!["address1", "email1", "id1", "name1"]
+            );
+        }
+    }
 }

crates/pgt_completions/src/sanitization.rs

@@ -1,4 +1,4 @@
-use std::borrow::Cow;
+use std::{borrow::Cow, cmp::max};
 
 use pgt_text_size::TextSize;
 
@@ -24,6 +24,7 @@ where
         if cursor_inbetween_nodes(params.tree, params.position)
             || cursor_prepared_to_write_token_after_last_node(params.tree, params.position)
             || cursor_before_semicolon(params.tree, params.position)
+            || cursor_on_a_dot(&params.text, params.position)
         {
             SanitizedCompletionParams::with_adjusted_sql(params)
         } else {
@@ -44,12 +45,13 @@ where
 
         let mut sql_iter = params.text.chars();
 
-        for idx in 0..cursor_pos + 1 {
+        let max = max(cursor_pos + 1, params.text.len());
+
+        for idx in 0..max {
             match sql_iter.next() {
                 Some(c) => {
                     if idx == cursor_pos {
                         sql.push_str(SANITIZED_TOKEN);
-                        sql.push(' ');
                     }
                     sql.push(c);
                 }
@@ -149,6 +151,11 @@ fn cursor_prepared_to_write_token_after_last_node(
     cursor_pos == tree.root_node().end_byte() + 1
 }
 
+fn cursor_on_a_dot(sql: &str, position: TextSize) -> bool {
+    let position: usize = position.into();
+    sql.chars().nth(position - 1).is_some_and(|c| c == '.')
+}
+
 fn cursor_before_semicolon(tree: &tree_sitter::Tree, position: TextSize) -> bool {
     let mut cursor = tree.walk();
     let mut leaf_node = tree.root_node();
@@ -198,7 +205,7 @@ mod tests {
     use pgt_text_size::TextSize;
 
     use crate::sanitization::{
-        cursor_before_semicolon, cursor_inbetween_nodes,
+        cursor_before_semicolon, cursor_inbetween_nodes, cursor_on_a_dot,
         cursor_prepared_to_write_token_after_last_node,
     };
 
@@ -263,6 +270,20 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn on_a_dot() {
+        let input = "select * from private.";
+
+        // select * from private.| <-- on a dot
+        assert!(cursor_on_a_dot(input, TextSize::new(22)));
+
+        // select * from private|. <-- before the dot
+        assert!(!cursor_on_a_dot(input, TextSize::new(21)));
+
+        // select * from private. | <-- too far off the dot
+        assert!(!cursor_on_a_dot(input, TextSize::new(23)));
+    }
+
     #[test]
     fn test_cursor_before_semicolon() {
         // Idx "13" is the exlusive end of `select * from` (first space after from)