fix bug

Author: juleswritescode

crates/pgt_completions/src/providers/columns.rs

@@ -260,4 +260,69 @@ mod tests {
                 .any(|item| item.kind == CompletionItemKind::Column)
         );
     }
+
+    #[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;
 
@@ -45,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);
                 }

crates/pgt_treesitter_queries/src/lib.rs

@@ -185,4 +185,31 @@ on sq1.id = pt.id;
         assert_eq!(results[0].get_schema(sql), Some("private".into()));
         assert_eq!(results[0].get_table(sql), "something");
     }
+
+    #[test]
+    fn picks_it_up_without_schemas() {
+        let sql = r#"select * from users"#;
+
+        let mut parser = tree_sitter::Parser::new();
+        parser.set_language(tree_sitter_sql::language()).unwrap();
+
+        let tree = parser.parse(sql, None).unwrap();
+
+        // trust me bro
+
+        let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), sql);
+
+        executor.add_query_results::<RelationMatch>();
+
+        let results: Vec<&RelationMatch> = executor
+            .get_iter(None)
+            .filter_map(|q| q.try_into().ok())
+            .collect();
+
+        println!("{:?}", results);
+
+        assert_eq!(results.len(), 1);
+        assert_eq!(results[0].get_schema(sql), None);
+        assert_eq!(results[0].get_table(sql), "users");
+    }
 }