Skip to content

Commit d022e2e

Browse files
juleswritescodejuleswritescode
committed
alter table statements
1 parent 66df5c6 commit d022e2e

File tree

3 files changed

+127
-23
lines changed

3 files changed

+127
-23
lines changed

‎crates/pgt_completions/src/context.rs

Lines changed: 37 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ pub enum WrappingClause<'a> {
2323
Delete,
2424
ColumnDefinitions,
2525
Insert,
26+
AlterTable,
27+
DropTable,
2628
}
2729

2830
#[derive(PartialEq, Eq, Debug)]
@@ -118,9 +120,6 @@ pub(crate) struct CompletionContext<'a> {
118120
pub is_invocation: bool,
119121
pub wrapping_statement_range: Option<tree_sitter::Range>,
120122

121-
/// Some incomplete statements can't be correctly parsed by TreeSitter.
122-
pub is_in_error_node: bool,
123-
124123
pub mentioned_relations: HashMap<Option<String>, HashSet<String>>,
125124
pub mentioned_table_aliases: HashMap<String, String>,
126125
pub mentioned_columns: HashMap<Option<WrappingClause<'a>>, HashSet<MentionedColumn>>,
@@ -142,12 +141,19 @@ impl<'a> CompletionContext<'a> {
142141
mentioned_relations: HashMap::new(),
143142
mentioned_table_aliases: HashMap::new(),
144143
mentioned_columns: HashMap::new(),
145-
is_in_error_node: false,
146144
};
147145

148146
ctx.gather_tree_context();
149147
ctx.gather_info_from_ts_queries();
150148

149+
// if cfg!(test) {
150+
// println!("{:?}", ctx.position);
151+
// println!("{:?}", ctx.text);
152+
// println!("{:?}", ctx.wrapping_clause_type);
153+
// println!("{:?}", ctx.wrapping_node_kind);
154+
// println!("{:?}", ctx.before_cursor_matches_kind(&["keyword_table"]));
155+
// }
156+
151157
ctx
152158
}
153159

@@ -284,7 +290,7 @@ impl<'a> CompletionContext<'a> {
284290
}
285291

286292
// try to gather context from the siblings if we're within an error node.
287-
if self.is_in_error_node {
293+
if parent_node_kind == "ERROR" {
288294
if let Some(clause_type) = self.get_wrapping_clause_from_siblings(current_node) {
289295
self.wrapping_clause_type = Some(clause_type);
290296
}
@@ -309,7 +315,8 @@ impl<'a> CompletionContext<'a> {
309315
}
310316
}
311317

312-
"where" | "update" | "select" | "delete" | "from" | "join" | "column_definitions" => {
318+
"where" | "update" | "select" | "delete" | "from" | "join" | "column_definitions"
319+
| "drop_table" | "alter_table" => {
313320
self.wrapping_clause_type =
314321
self.get_wrapping_clause_from_current_node(current_node, &mut cursor);
315322
}
@@ -318,10 +325,6 @@ impl<'a> CompletionContext<'a> {
318325
self.wrapping_node_kind = current_node_kind.try_into().ok();
319326
}
320327

321-
"ERROR" => {
322-
self.is_in_error_node = true;
323-
}
324-
325328
_ => {}
326329
}
327330

@@ -372,6 +375,16 @@ impl<'a> CompletionContext<'a> {
372375
(WrappingClause::Insert, &["insert", "into"]),
373376
(WrappingClause::From, &["from"]),
374377
(WrappingClause::Join { on_node: None }, &["join"]),
378+
(WrappingClause::AlterTable, &["alter", "table"]),
379+
(
380+
WrappingClause::AlterTable,
381+
&["alter", "table", "if", "exists"],
382+
),
383+
(WrappingClause::DropTable, &["drop", "table"]),
384+
(
385+
WrappingClause::DropTable,
386+
&["drop", "table", "if", "exists"],
387+
),
375388
];
376389

377390
let first_sibling = self.get_first_sibling(node);
@@ -431,6 +444,8 @@ impl<'a> CompletionContext<'a> {
431444
"select" => Some(WrappingClause::Select),
432445
"delete" => Some(WrappingClause::Delete),
433446
"from" => Some(WrappingClause::From),
447+
"drop_table" => Some(WrappingClause::DropTable),
448+
"alter_table" => Some(WrappingClause::AlterTable),
434449
"column_definitions" => Some(WrappingClause::ColumnDefinitions),
435450
"insert" => Some(WrappingClause::Insert),
436451
"join" => {
@@ -449,6 +464,18 @@ impl<'a> CompletionContext<'a> {
449464
_ => None,
450465
}
451466
}
467+
468+
pub(crate) fn before_cursor_matches_kind(&self, kinds: &[&'static str]) -> bool {
469+
self.node_under_cursor.is_some_and(|mut node| {
470+
// move up to the parent until we're at top OR we have a prev sibling
471+
while node.prev_sibling().is_none() && node.parent().is_some() {
472+
node = node.parent().unwrap();
473+
}
474+
475+
node.prev_sibling()
476+
.is_some_and(|sib| kinds.contains(&sib.kind()))
477+
})
478+
}
452479
}
453480

454481
#[cfg(test)]

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -310,4 +310,73 @@ mod tests {
310310
)
311311
.await;
312312
}
313+
314+
#[tokio::test]
315+
async fn suggests_tables_in_alter_and_drop_statements() {
316+
let setup = r#"
317+
create schema auth;
318+
319+
create table auth.users (
320+
uid serial primary key,
321+
name text not null,
322+
email text unique not null
323+
);
324+
325+
create table auth.posts (
326+
pid serial primary key,
327+
user_id int not null references auth.users(uid),
328+
title text not null,
329+
content text,
330+
created_at timestamp default now()
331+
);
332+
"#;
333+
334+
assert_complete_results(
335+
format!("alter table {}", CURSOR_POS).as_str(),
336+
vec![
337+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
338+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
339+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
340+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
341+
],
342+
setup,
343+
)
344+
.await;
345+
346+
assert_complete_results(
347+
format!("alter table if exists {}", CURSOR_POS).as_str(),
348+
vec![
349+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
350+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
351+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
352+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
353+
],
354+
setup,
355+
)
356+
.await;
357+
358+
assert_complete_results(
359+
format!("drop table {}", CURSOR_POS).as_str(),
360+
vec![
361+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
362+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
363+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
364+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
365+
],
366+
setup,
367+
)
368+
.await;
369+
370+
assert_complete_results(
371+
format!("drop table if exists {}", CURSOR_POS).as_str(),
372+
vec![
373+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
374+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
375+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join
376+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
377+
],
378+
setup,
379+
)
380+
.await;
381+
}
313382
}

‎crates/pgt_completions/src/relevance/filtering.rs

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,7 @@ impl CompletionFilter<'_> {
2424
}
2525

2626
fn completable_context(&self, ctx: &CompletionContext) -> Option<()> {
27-
if ctx.wrapping_node_kind.is_none()
28-
&& ctx.wrapping_clause_type.is_none()
29-
&& ctx.is_in_error_node
30-
{
27+
if ctx.wrapping_node_kind.is_none() && ctx.wrapping_clause_type.is_none() {
3128
return None;
3229
}
3330

@@ -78,17 +75,24 @@ impl CompletionFilter<'_> {
7875
ctx.wrapping_node_kind
7976
.as_ref()
8077
.is_some_and(|n| n != &WrappingNode::List)
81-
&& ctx.node_under_cursor.is_some_and(|n| {
82-
n.prev_sibling()
83-
.is_some_and(|sib| sib.kind() == "keyword_into")
84-
})
78+
&& ctx.before_cursor_matches_kind(&["keyword_into"])
8579
}
8680

81+
WrappingClause::DropTable | WrappingClause::AlterTable => ctx
82+
.before_cursor_matches_kind(&[
83+
"keyword_exists",
84+
"keyword_only",
85+
"keyword_table",
86+
]),
87+
8788
_ => true,
8889
},
8990
CompletionRelevanceData::Column(_) => {
9091
match clause {
91-
WrappingClause::From | WrappingClause::ColumnDefinitions => false,
92+
WrappingClause::From
93+
| WrappingClause::ColumnDefinitions
94+
| WrappingClause::AlterTable
95+
| WrappingClause::DropTable => false,
9296

9397
// We can complete columns in JOIN cluases, but only if we are after the
9498
// ON node in the "ON u.id = posts.user_id" part.
@@ -123,14 +127,18 @@ impl CompletionFilter<'_> {
123127
| WrappingClause::Update
124128
| WrappingClause::Delete => true,
125129

130+
WrappingClause::DropTable | WrappingClause::AlterTable => ctx
131+
.before_cursor_matches_kind(&[
132+
"keyword_exists",
133+
"keyword_only",
134+
"keyword_table",
135+
]),
136+
126137
WrappingClause::Insert => {
127138
ctx.wrapping_node_kind
128139
.as_ref()
129140
.is_some_and(|n| n != &WrappingNode::List)
130-
&& ctx.node_under_cursor.is_some_and(|n| {
131-
n.prev_sibling()
132-
.is_some_and(|sib| sib.kind() == "keyword_into")
133-
})
141+
&& ctx.before_cursor_matches_kind(&["keyword_into"])
134142
}
135143

136144
WrappingClause::ColumnDefinitions => false,

0 commit comments

Comments
 (0)