Skip to content

Commit 513c73c

Browse files
juleswritescodejuleswritescode
authored
feat(completions): alter/drop/rename column statements (#421)
1 parent d7ad288 commit 513c73c

File tree

4 files changed

+191
-32
lines changed

4 files changed

+191
-32
lines changed

‎crates/pgt_completions/src/context/mod.rs

Lines changed: 60 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ pub enum WrappingClause<'a> {
3131
Insert,
3232
AlterTable,
3333
DropTable,
34+
DropColumn,
35+
AlterColumn,
36+
RenameColumn,
3437
PolicyName,
3538
ToRoleAssignment,
3639
}
@@ -424,7 +427,7 @@ impl<'a> CompletionContext<'a> {
424427
}
425428

426429
"where" | "update" | "select" | "delete" | "from" | "join" | "column_definitions"
427-
| "drop_table" | "alter_table" => {
430+
| "drop_table" | "alter_table" | "drop_column" | "alter_column" | "rename_column" => {
428431
self.wrapping_clause_type =
429432
self.get_wrapping_clause_from_current_node(current_node, &mut cursor);
430433
}
@@ -515,6 +518,8 @@ impl<'a> CompletionContext<'a> {
515518
(WrappingClause::From, &["from"]),
516519
(WrappingClause::Join { on_node: None }, &["join"]),
517520
(WrappingClause::AlterTable, &["alter", "table"]),
521+
(WrappingClause::AlterColumn, &["alter", "table", "alter"]),
522+
(WrappingClause::RenameColumn, &["alter", "table", "rename"]),
518523
(
519524
WrappingClause::AlterTable,
520525
&["alter", "table", "if", "exists"],
@@ -575,10 +580,54 @@ impl<'a> CompletionContext<'a> {
575580
let mut first_sibling = self.get_first_sibling(node);
576581

577582
if let Some(clause) = self.wrapping_clause_type.as_ref() {
578-
if clause == &WrappingClause::Insert {
579-
while let Some(sib) = first_sibling.next_sibling() {
580-
match sib.kind() {
581-
"object_reference" => {
583+
match *clause {
584+
WrappingClause::Insert => {
585+
while let Some(sib) = first_sibling.next_sibling() {
586+
match sib.kind() {
587+
"object_reference" => {
588+
if let Some(NodeText::Original(txt)) =
589+
self.get_ts_node_content(&sib)
590+
{
591+
let mut iter = txt.split('.').rev();
592+
let table = iter.next().unwrap().to_string();
593+
let schema = iter.next().map(|s| s.to_string());
594+
self.mentioned_relations
595+
.entry(schema)
596+
.and_modify(|s| {
597+
s.insert(table.clone());
598+
})
599+
.or_insert(HashSet::from([table]));
600+
}
601+
}
602+
603+
"column" => {
604+
if let Some(NodeText::Original(txt)) =
605+
self.get_ts_node_content(&sib)
606+
{
607+
let entry = MentionedColumn {
608+
column: txt,
609+
alias: None,
610+
};
611+
612+
self.mentioned_columns
613+
.entry(Some(WrappingClause::Insert))
614+
.and_modify(|s| {
615+
s.insert(entry.clone());
616+
})
617+
.or_insert(HashSet::from([entry]));
618+
}
619+
}
620+
621+
_ => {}
622+
}
623+
624+
first_sibling = sib;
625+
}
626+
}
627+
628+
WrappingClause::AlterColumn => {
629+
while let Some(sib) = first_sibling.next_sibling() {
630+
if sib.kind() == "object_reference" {
582631
if let Some(NodeText::Original(txt)) = self.get_ts_node_content(&sib) {
583632
let mut iter = txt.split('.').rev();
584633
let table = iter.next().unwrap().to_string();
@@ -591,27 +640,12 @@ impl<'a> CompletionContext<'a> {
591640
.or_insert(HashSet::from([table]));
592641
}
593642
}
594-
"column" => {
595-
if let Some(NodeText::Original(txt)) = self.get_ts_node_content(&sib) {
596-
let entry = MentionedColumn {
597-
column: txt,
598-
alias: None,
599-
};
600643

601-
self.mentioned_columns
602-
.entry(Some(WrappingClause::Insert))
603-
.and_modify(|s| {
604-
s.insert(entry.clone());
605-
})
606-
.or_insert(HashSet::from([entry]));
607-
}
608-
}
609-
610-
_ => {}
644+
first_sibling = sib;
611645
}
612-
613-
first_sibling = sib;
614646
}
647+
648+
_ => {}
615649
}
616650
}
617651
}
@@ -628,6 +662,9 @@ impl<'a> CompletionContext<'a> {
628662
"delete" => Some(WrappingClause::Delete),
629663
"from" => Some(WrappingClause::From),
630664
"drop_table" => Some(WrappingClause::DropTable),
665+
"drop_column" => Some(WrappingClause::DropColumn),
666+
"alter_column" => Some(WrappingClause::AlterColumn),
667+
"rename_column" => Some(WrappingClause::RenameColumn),
631668
"alter_table" => Some(WrappingClause::AlterTable),
632669
"column_definitions" => Some(WrappingClause::ColumnDefinitions),
633670
"insert" => Some(WrappingClause::Insert),

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,4 +762,59 @@ mod tests {
762762
)
763763
.await;
764764
}
765+
766+
#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
767+
async fn suggests_columns_in_alter_table_and_drop_table(pool: PgPool) {
768+
let setup = r#"
769+
create table instruments (
770+
id bigint primary key generated always as identity,
771+
name text not null,
772+
z text,
773+
created_at timestamp with time zone default now()
774+
);
775+
776+
create table others (
777+
a text,
778+
b text,
779+
c text
780+
);
781+
"#;
782+
783+
pool.execute(setup).await.unwrap();
784+
785+
let queries = vec![
786+
format!("alter table instruments drop column {}", CURSOR_POS),
787+
format!(
788+
"alter table instruments drop column if exists {}",
789+
CURSOR_POS
790+
),
791+
format!(
792+
"alter table instruments alter column {} set default",
793+
CURSOR_POS
794+
),
795+
format!("alter table instruments alter {} set default", CURSOR_POS),
796+
format!("alter table public.instruments alter column {}", CURSOR_POS),
797+
format!("alter table instruments alter {}", CURSOR_POS),
798+
format!("alter table instruments rename {} to new_col", CURSOR_POS),
799+
format!(
800+
"alter table public.instruments rename column {} to new_col",
801+
CURSOR_POS
802+
),
803+
];
804+
805+
for query in queries {
806+
assert_complete_results(
807+
query.as_str(),
808+
vec![
809+
CompletionAssertion::Label("created_at".into()),
810+
CompletionAssertion::Label("id".into()),
811+
CompletionAssertion::Label("name".into()),
812+
CompletionAssertion::Label("z".into()),
813+
],
814+
None,
815+
&pool,
816+
)
817+
.await;
818+
}
819+
}
765820
}

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

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,13 @@ impl CompletionFilter<'_> {
7474
.map(|clause| {
7575
match self.data {
7676
CompletionRelevanceData::Table(_) => match clause {
77-
WrappingClause::Select
78-
| WrappingClause::Where
79-
| WrappingClause::ColumnDefinitions => false,
77+
WrappingClause::From | WrappingClause::Update => true,
78+
79+
WrappingClause::Join { on_node: None } => true,
80+
WrappingClause::Join { on_node: Some(on) } => ctx
81+
.node_under_cursor
82+
.as_ref()
83+
.is_some_and(|cn| cn.start_byte() < on.end_byte()),
8084

8185
WrappingClause::Insert => {
8286
ctx.wrapping_node_kind
@@ -94,15 +98,22 @@ impl CompletionFilter<'_> {
9498
"keyword_table",
9599
]),
96100

97-
_ => true,
101+
_ => false,
98102
},
99103

100104
CompletionRelevanceData::Column(_) => {
101105
match clause {
102-
WrappingClause::From
103-
| WrappingClause::ColumnDefinitions
104-
| WrappingClause::AlterTable
105-
| WrappingClause::DropTable => false,
106+
WrappingClause::Select
107+
| WrappingClause::Update
108+
| WrappingClause::Delete
109+
| WrappingClause::DropColumn => true,
110+
111+
WrappingClause::RenameColumn => ctx
112+
.before_cursor_matches_kind(&["keyword_rename", "keyword_column"]),
113+
114+
WrappingClause::AlterColumn => {
115+
ctx.before_cursor_matches_kind(&["keyword_alter", "keyword_column"])
116+
}
106117

107118
// We can complete columns in JOIN cluases, but only if we are after the
108119
// ON node in the "ON u.id = posts.user_id" part.
@@ -126,7 +137,7 @@ impl CompletionFilter<'_> {
126137
&& ctx.parent_matches_one_of_kind(&["field"]))
127138
}
128139

129-
_ => true,
140+
_ => false,
130141
}
131142
}
132143

‎crates/pgt_treesitter_queries/src/queries/relations.rs

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,16 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
2222
(identifier)? @table
2323
)+
2424
)
25+
(alter_table
26+
(keyword_alter)
27+
(keyword_table)
28+
(object_reference
29+
.
30+
(identifier) @schema_or_table
31+
"."?
32+
(identifier)? @table
33+
)+
34+
)
2535
"#;
2636
tree_sitter::Query::new(tree_sitter_sql::language(), QUERY_STR).expect("Invalid TS Query")
2737
});
@@ -196,4 +206,50 @@ mod tests {
196206
assert_eq!(results[0].get_schema(sql), None);
197207
assert_eq!(results[0].get_table(sql), "users");
198208
}
209+
210+
#[test]
211+
fn finds_alter_table_with_schema() {
212+
let sql = r#"alter table public.users alter some_col set default 15;"#;
213+
214+
let mut parser = tree_sitter::Parser::new();
215+
parser.set_language(tree_sitter_sql::language()).unwrap();
216+
217+
let tree = parser.parse(sql, None).unwrap();
218+
219+
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), sql);
220+
221+
executor.add_query_results::<RelationMatch>();
222+
223+
let results: Vec<&RelationMatch> = executor
224+
.get_iter(None)
225+
.filter_map(|q| q.try_into().ok())
226+
.collect();
227+
228+
assert_eq!(results.len(), 1);
229+
assert_eq!(results[0].get_schema(sql), Some("public".into()));
230+
assert_eq!(results[0].get_table(sql), "users");
231+
}
232+
233+
#[test]
234+
fn finds_alter_table_without_schema() {
235+
let sql = r#"alter table users alter some_col set default 15;"#;
236+
237+
let mut parser = tree_sitter::Parser::new();
238+
parser.set_language(tree_sitter_sql::language()).unwrap();
239+
240+
let tree = parser.parse(sql, None).unwrap();
241+
242+
let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), sql);
243+
244+
executor.add_query_results::<RelationMatch>();
245+
246+
let results: Vec<&RelationMatch> = executor
247+
.get_iter(None)
248+
.filter_map(|q| q.try_into().ok())
249+
.collect();
250+
251+
assert_eq!(results.len(), 1);
252+
assert_eq!(results[0].get_schema(sql), None);
253+
assert_eq!(results[0].get_table(sql), "users");
254+
}
199255
}

0 commit comments

Comments
 (0)