Skip to content

Commit 4ee920d

Browse files
juleswritescodejuleswritescode
committed
beautiful
1 parent 089c50e commit 4ee920d

File tree

9 files changed

+245
-98
lines changed

9 files changed

+245
-98
lines changed

‎crates/pg_completions/src/context.rs

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,4 @@
1-
use std::{
2-
collections::{HashMap, HashSet},
3-
hash::Hash,
4-
ops::Range,
5-
};
1+
use std::collections::{HashMap, HashSet};
62

73
use pg_schema_cache::SchemaCache;
84
use pg_treesitter_queries::{
@@ -62,7 +58,7 @@ pub(crate) struct CompletionContext<'a> {
6258
pub schema_name: Option<String>,
6359
pub wrapping_clause_type: Option<ClauseType>,
6460
pub is_invocation: bool,
65-
pub wrapping_statement_range: Option<Range<usize>>,
61+
pub wrapping_statement_range: Option<tree_sitter::Range>,
6662

6763
pub mentioned_relations: HashMap<Option<String>, HashSet<String>>,
6864
}
@@ -74,7 +70,6 @@ impl<'a> CompletionContext<'a> {
7470
text: &params.text,
7571
schema_cache: params.schema,
7672
position: usize::from(params.position),
77-
7873
ts_node: None,
7974
schema_name: None,
8075
wrapping_clause_type: None,
@@ -86,6 +81,8 @@ impl<'a> CompletionContext<'a> {
8681
ctx.gather_tree_context();
8782
ctx.gather_info_from_ts_queries();
8883

84+
dbg!(ctx.wrapping_statement_range);
85+
8986
ctx
9087
}
9188

@@ -98,6 +95,8 @@ impl<'a> CompletionContext<'a> {
9895
let stmt_range = self.wrapping_statement_range.as_ref();
9996
let sql = self.text;
10097

98+
dbg!(sql);
99+
101100
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), self.text);
102101

103102
executor.add_query_results::<queries::RelationMatch>();
@@ -174,9 +173,9 @@ impl<'a> CompletionContext<'a> {
174173
}
175174

176175
match previous_node.kind() {
177-
"statement" => {
176+
"statement" | "subquery" => {
178177
self.wrapping_clause_type = current_node.kind().try_into().ok();
179-
self.wrapping_statement_range = Some(previous_node.byte_range());
178+
self.wrapping_statement_range = Some(previous_node.range());
180179
}
181180
"invocation" => self.is_invocation = true,
182181

@@ -263,7 +262,7 @@ mod tests {
263262
];
264263

265264
for (query, expected_clause) in test_cases {
266-
let (position, text) = get_text_and_position(query.as_str());
265+
let (position, text) = get_text_and_position(query.as_str().into());
267266

268267
let tree = get_tree(text.as_str());
269268

@@ -296,7 +295,7 @@ mod tests {
296295
];
297296

298297
for (query, expected_schema) in test_cases {
299-
let (position, text) = get_text_and_position(query.as_str());
298+
let (position, text) = get_text_and_position(query.as_str().into());
300299

301300
let tree = get_tree(text.as_str());
302301
let params = crate::CompletionParams {
@@ -330,7 +329,7 @@ mod tests {
330329
];
331330

332331
for (query, is_invocation) in test_cases {
333-
let (position, text) = get_text_and_position(query.as_str());
332+
let (position, text) = get_text_and_position(query.as_str().into());
334333

335334
let tree = get_tree(text.as_str());
336335
let params = crate::CompletionParams {
@@ -354,7 +353,7 @@ mod tests {
354353
];
355354

356355
for query in cases {
357-
let (position, text) = get_text_and_position(query.as_str());
356+
let (position, text) = get_text_and_position(query.as_str().into());
358357

359358
let tree = get_tree(text.as_str());
360359

@@ -382,7 +381,7 @@ mod tests {
382381
fn does_not_fail_on_trailing_whitespace() {
383382
let query = format!("select * from {}", CURSOR_POS);
384383

385-
let (position, text) = get_text_and_position(query.as_str());
384+
let (position, text) = get_text_and_position(query.as_str().into());
386385

387386
let tree = get_tree(text.as_str());
388387

@@ -408,7 +407,7 @@ mod tests {
408407
fn does_not_fail_with_empty_statements() {
409408
let query = format!("{}", CURSOR_POS);
410409

411-
let (position, text) = get_text_and_position(query.as_str());
410+
let (position, text) = get_text_and_position(query.as_str().into());
412411

413412
let tree = get_tree(text.as_str());
414413

@@ -433,7 +432,7 @@ mod tests {
433432
// is selecting a certain column name, such as `frozen_account`.
434433
let query = format!("select * fro{}", CURSOR_POS);
435434

436-
let (position, text) = get_text_and_position(query.as_str());
435+
let (position, text) = get_text_and_position(query.as_str().into());
437436

438437
let tree = get_tree(text.as_str());
439438

‎crates/pg_completions/src/providers/columns.rs

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,97 @@ pub fn complete_columns(ctx: &CompletionContext, builder: &mut CompletionBuilder
1818
builder.add_item(item);
1919
}
2020
}
21+
22+
#[cfg(test)]
23+
mod tests {
24+
use crate::{
25+
complete,
26+
test_helper::{get_test_deps, get_test_params, InputQuery, CURSOR_POS},
27+
CompletionItem,
28+
};
29+
30+
struct TestCase {
31+
query: String,
32+
message: &'static str,
33+
label: &'static str,
34+
description: &'static str,
35+
}
36+
37+
impl TestCase {
38+
fn get_input_query(&self) -> InputQuery {
39+
let strs: Vec<&str> = self.query.split_whitespace().collect();
40+
strs.join(" ").as_str().into()
41+
}
42+
}
43+
44+
#[tokio::test]
45+
async fn completes_columns() {
46+
let setup = r#"
47+
create schema private;
48+
49+
create table public.users (
50+
id serial primary key,
51+
name text
52+
);
53+
54+
create table public.audio_books (
55+
id serial primary key,
56+
narrator text
57+
);
58+
59+
create table private.audio_books (
60+
id serial primary key,
61+
narrator_id text
62+
);
63+
"#;
64+
65+
let queries: Vec<TestCase> = vec![
66+
TestCase {
67+
message: "correctly prefers the columns of present tables",
68+
query: format!(r#"select na{} from public.audio_books;"#, CURSOR_POS),
69+
label: "narrator",
70+
description: "Table: public.audio_books",
71+
},
72+
TestCase {
73+
message: "correctly handles nested queries",
74+
query: format!(
75+
r#"
76+
select
77+
*
78+
from (
79+
select id, na{}
80+
from private.audio_books
81+
) as subquery
82+
join public.users u
83+
on u.id = subquery.id;
84+
"#,
85+
CURSOR_POS
86+
),
87+
label: "narrator_id",
88+
description: "Table: private.audio_books",
89+
},
90+
TestCase {
91+
message: "works without a schema",
92+
query: format!(r#"select na{} from users;"#, CURSOR_POS),
93+
label: "name",
94+
description: "Table: public.users",
95+
},
96+
];
97+
98+
for q in queries {
99+
let (tree, cache) = get_test_deps(setup, q.get_input_query()).await;
100+
let params = get_test_params(&tree, &cache, q.get_input_query());
101+
let results = complete(params);
102+
103+
let CompletionItem {
104+
label, description, ..
105+
} = results
106+
.into_iter()
107+
.next()
108+
.expect("Should return at least one completion item");
109+
110+
assert_eq!(label, q.label, "{}", q.message);
111+
assert_eq!(description, q.description, "{}", q.message);
112+
}
113+
}
114+
}

‎crates/pg_completions/src/providers/functions.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,8 @@ mod tests {
4343

4444
let query = format!("select coo{}", CURSOR_POS);
4545

46-
let (tree, cache) = get_test_deps(setup, &query).await;
47-
let params = get_test_params(&tree, &cache, &query);
46+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
47+
let params = get_test_params(&tree, &cache, query.as_str().into());
4848
let results = complete(params);
4949

5050
let CompletionItem { label, .. } = results
@@ -76,8 +76,8 @@ mod tests {
7676

7777
let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
7878

79-
let (tree, cache) = get_test_deps(setup, &query).await;
80-
let params = get_test_params(&tree, &cache, &query);
79+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
80+
let params = get_test_params(&tree, &cache, query.as_str().into());
8181
let results = complete(params);
8282

8383
let CompletionItem { label, kind, .. } = results
@@ -110,8 +110,8 @@ mod tests {
110110

111111
let query = format!(r#"select coo{}"#, CURSOR_POS);
112112

113-
let (tree, cache) = get_test_deps(setup, &query).await;
114-
let params = get_test_params(&tree, &cache, &query);
113+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
114+
let params = get_test_params(&tree, &cache, query.as_str().into());
115115
let results = complete(params);
116116

117117
let CompletionItem { label, kind, .. } = results
@@ -144,8 +144,8 @@ mod tests {
144144

145145
let query = format!(r#"select * from coo{}()"#, CURSOR_POS);
146146

147-
let (tree, cache) = get_test_deps(setup, &query).await;
148-
let params = get_test_params(&tree, &cache, &query);
147+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
148+
let params = get_test_params(&tree, &cache, query.as_str().into());
149149
let results = complete(params);
150150

151151
let CompletionItem { label, kind, .. } = results

‎crates/pg_completions/src/providers/tables.rs

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ pub fn complete_tables(ctx: &CompletionContext, builder: &mut CompletionBuilder)
2323

2424
#[cfg(test)]
2525
mod tests {
26+
2627
use crate::{
2728
complete,
2829
test_helper::{get_test_deps, get_test_params, CURSOR_POS},
@@ -41,8 +42,8 @@ mod tests {
4142

4243
let query = format!("select * from u{}", CURSOR_POS);
4344

44-
let (tree, cache) = get_test_deps(setup, &query).await;
45-
let params = get_test_params(&tree, &cache, &query);
45+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
46+
let params = get_test_params(&tree, &cache, query.as_str().into());
4647
let results = complete(params);
4748

4849
assert!(!results.items.is_empty());
@@ -79,8 +80,8 @@ mod tests {
7980
];
8081

8182
for (query, expected_label) in test_cases {
82-
let (tree, cache) = get_test_deps(setup, &query).await;
83-
let params = get_test_params(&tree, &cache, &query);
83+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
84+
let params = get_test_params(&tree, &cache, query.as_str().into());
8485
let results = complete(params);
8586

8687
assert!(!results.items.is_empty());
@@ -124,8 +125,8 @@ mod tests {
124125
];
125126

126127
for (query, expected_label) in test_cases {
127-
let (tree, cache) = get_test_deps(setup, &query).await;
128-
let params = get_test_params(&tree, &cache, &query);
128+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
129+
let params = get_test_params(&tree, &cache, query.as_str().into());
129130
let results = complete(params);
130131

131132
assert!(!results.items.is_empty());
@@ -161,8 +162,8 @@ mod tests {
161162

162163
let query = format!(r#"select * from coo{}"#, CURSOR_POS);
163164

164-
let (tree, cache) = get_test_deps(setup, &query).await;
165-
let params = get_test_params(&tree, &cache, &query);
165+
let (tree, cache) = get_test_deps(setup, query.as_str().into()).await;
166+
let params = get_test_params(&tree, &cache, query.as_str().into());
166167
let results = complete(params);
167168

168169
let CompletionItem { label, kind, .. } = results

‎crates/pg_completions/src/test_helper.rs

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,34 @@ use crate::CompletionParams;
66

77
pub static CURSOR_POS: char = '€';
88

9+
pub struct InputQuery {
10+
sql: String,
11+
position: usize,
12+
}
13+
14+
impl From<&str> for InputQuery {
15+
fn from(value: &str) -> Self {
16+
let position = value
17+
.find(CURSOR_POS)
18+
.map(|p| p.saturating_sub(1))
19+
.expect("Insert Cursor Position into your Query.");
20+
21+
InputQuery {
22+
sql: value.replace(CURSOR_POS, ""),
23+
position,
24+
}
25+
}
26+
}
27+
28+
impl ToString for InputQuery {
29+
fn to_string(&self) -> String {
30+
self.sql.clone()
31+
}
32+
}
33+
934
pub(crate) async fn get_test_deps(
1035
setup: &str,
11-
input: &str,
36+
input: InputQuery,
1237
) -> (tree_sitter::Tree, pg_schema_cache::SchemaCache) {
1338
let test_db = get_new_test_db().await;
1439

@@ -26,27 +51,19 @@ pub(crate) async fn get_test_deps(
2651
.set_language(tree_sitter_sql::language())
2752
.expect("Error loading sql language");
2853

29-
let tree = parser.parse(input, None).unwrap();
54+
let tree = parser.parse(&input.to_string(), None).unwrap();
3055

3156
(tree, schema_cache)
3257
}
3358

34-
pub(crate) fn get_text_and_position(sql: &str) -> (usize, String) {
35-
// the cursor is to the left of the `CURSOR_POS`
36-
let position = sql
37-
.find(CURSOR_POS)
38-
.expect("Please insert the CURSOR_POS into your query.")
39-
.saturating_sub(1);
40-
41-
let text = sql.replace(CURSOR_POS, "");
42-
43-
(position, text)
59+
pub(crate) fn get_text_and_position(q: InputQuery) -> (usize, String) {
60+
(q.position, q.sql)
4461
}
4562

4663
pub(crate) fn get_test_params<'a>(
4764
tree: &'a tree_sitter::Tree,
4865
schema_cache: &'a pg_schema_cache::SchemaCache,
49-
sql: &'a str,
66+
sql: InputQuery,
5067
) -> CompletionParams<'a> {
5168
let (position, text) = get_text_and_position(sql);
5269

‎crates/pg_lsp_new/src/handlers/completions.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,5 +56,6 @@ fn to_lsp_types_completion_item_kind(
5656
match pg_comp_kind {
5757
pg_completions::CompletionItemKind::Function
5858
| pg_completions::CompletionItemKind::Table => lsp_types::CompletionItemKind::CLASS,
59+
pg_completions::CompletionItemKind::Column => lsp_types::CompletionItemKind::FIELD,
5960
}
6061
}

0 commit comments

Comments
 (0)