Skip to content

Commit 3db26f4

Browse files
juleswritescodejuleswritescode
authored
refactor(completions): remove custom parse for policies, use tree-sitter (#530)
1 parent 21ce254 commit 3db26f4

File tree

15 files changed

+301
-701991
lines changed

15 files changed

+301
-701991
lines changed

‎.github/workflows/pull_request.yml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,9 @@ jobs:
112112
sleep 1
113113
done
114114
115+
- name: Install tree-sitter-cli
116+
run: cargo install tree-sitter-cli
117+
115118
- name: Setup sqlx-cli
116119
run: cargo install sqlx-cli
117120

@@ -167,6 +170,9 @@ jobs:
167170
- name: Setup Postgres
168171
uses: ./.github/actions/setup-postgres
169172

173+
- name: Install tree-sitter-cli
174+
run: cargo install tree-sitter-cli
175+
170176
- name: Run tests
171177
run: cargo test --workspace
172178

@@ -195,6 +201,8 @@ jobs:
195201
uses: moonrepo/setup-rust@v1
196202
with:
197203
cache-base: main
204+
- name: Install tree-sitter-cli
205+
run: cargo install tree-sitter-cli
198206
- name: Build main binary
199207
run: cargo build -p pgt_cli --release
200208
- name: Setup Bun
@@ -236,6 +244,8 @@ jobs:
236244
cache-base: main
237245
env:
238246
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
247+
- name: Install tree-sitter-cli
248+
run: cargo install tree-sitter-cli
239249
- name: Ensure RustFMT on nightly toolchain
240250
run: rustup component add rustfmt --toolchain nightly
241251
- name: echo toolchain

‎.gitignore‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,3 +26,8 @@ node_modules/
2626
.claude-session-id
2727

2828
squawk/
29+
30+
# Auto generated treesitter files
31+
crates/pgt_treesitter_grammar/src/grammar.json
32+
crates/pgt_treesitter_grammar/src/node-types.json
33+
crates/pgt_treesitter_grammar/src/parser.c

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,14 +274,16 @@ mod tests {
274274
pool.execute(setup).await.unwrap();
275275

276276
let query = format!(
277-
r#"create policy "my_pol" on public.instruments for insert with check (id = {})"#,
277+
r#"create policy "my_pol" on public.coos for insert with check (id = {})"#,
278278
QueryWithCursorPosition::cursor_marker()
279279
);
280280

281281
assert_complete_results(
282282
query.as_str(),
283283
vec![
284284
CompletionAssertion::LabelNotExists("string_concat".into()),
285+
CompletionAssertion::LabelAndKind("id".into(), CompletionItemKind::Column),
286+
CompletionAssertion::LabelAndKind("name".into(), CompletionItemKind::Column),
285287
CompletionAssertion::LabelAndKind(
286288
"my_cool_foo".into(),
287289
CompletionItemKind::Function,

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,21 @@ mod tests {
8080

8181
pool.execute(setup).await.unwrap();
8282

83+
assert_complete_results(
84+
format!(
85+
"alter policy \"{}",
86+
QueryWithCursorPosition::cursor_marker()
87+
)
88+
.as_str(),
89+
vec![
90+
CompletionAssertion::Label("read for public users disallowed".into()),
91+
CompletionAssertion::Label("write for public users allowed".into()),
92+
],
93+
None,
94+
&pool,
95+
)
96+
.await;
97+
8398
assert_complete_results(
8499
format!(
85100
"alter policy \"{}\" on private.users;",

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -656,4 +656,57 @@ mod tests {
656656
)
657657
.await;
658658
}
659+
660+
#[sqlx::test(migrator = "pgt_test_utils::MIGRATIONS")]
661+
async fn completes_tables_in_policies(pool: PgPool) {
662+
let setup = r#"
663+
create schema auth;
664+
665+
create table auth.users (
666+
uid serial primary key,
667+
name text not null,
668+
email text unique not null
669+
);
670+
671+
create table auth.posts (
672+
pid serial primary key,
673+
user_id int not null references auth.users(uid),
674+
title text not null,
675+
content text,
676+
created_at timestamp default now()
677+
);
678+
"#;
679+
680+
pool.execute(setup).await.unwrap();
681+
682+
assert_complete_results(
683+
format!(
684+
r#"create policy "my cool pol" on {}"#,
685+
QueryWithCursorPosition::cursor_marker()
686+
)
687+
.as_str(),
688+
vec![
689+
CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
690+
CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
691+
],
692+
None,
693+
&pool,
694+
)
695+
.await;
696+
697+
assert_complete_results(
698+
format!(
699+
r#"create policy "my cool pol" on auth.{}"#,
700+
QueryWithCursorPosition::cursor_marker()
701+
)
702+
.as_str(),
703+
vec![
704+
CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
705+
CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
706+
],
707+
None,
708+
&pool,
709+
)
710+
.await;
711+
}
659712
}

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

Lines changed: 35 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use pgt_schema_cache::ProcKind;
2-
32
use pgt_treesitter::context::{NodeUnderCursor, TreesitterContext, WrappingClause, WrappingNode};
43

54
use super::CompletionRelevanceData;
@@ -58,8 +57,7 @@ impl CompletionFilter<'_> {
5857
| Some(WrappingClause::Insert)
5958
| Some(WrappingClause::DropColumn)
6059
| Some(WrappingClause::AlterColumn)
61-
| Some(WrappingClause::RenameColumn)
62-
| Some(WrappingClause::PolicyCheck) => {
60+
| Some(WrappingClause::RenameColumn) => {
6361
// the literal is probably a column
6462
}
6563
_ => return None,
@@ -123,6 +121,13 @@ impl CompletionFilter<'_> {
123121
"keyword_table",
124122
]),
125123

124+
WrappingClause::CreatePolicy
125+
| WrappingClause::AlterPolicy
126+
| WrappingClause::DropPolicy => {
127+
ctx.matches_ancestor_history(&["object_reference"])
128+
&& ctx.before_cursor_matches_kind(&["keyword_on", "."])
129+
}
130+
126131
_ => false,
127132
},
128133

@@ -162,8 +167,11 @@ impl CompletionFilter<'_> {
162167
&& ctx.matches_ancestor_history(&["field"]))
163168
}
164169

165-
WrappingClause::PolicyCheck => {
166-
ctx.before_cursor_matches_kind(&["keyword_and", "("])
170+
WrappingClause::CheckOrUsingClause => {
171+
ctx.before_cursor_matches_kind(&["(", "keyword_and"])
172+
|| ctx.wrapping_node_kind.as_ref().is_some_and(|nk| {
173+
matches!(nk, WrappingNode::BinaryExpression)
174+
})
167175
}
168176

169177
_ => false,
@@ -176,9 +184,12 @@ impl CompletionFilter<'_> {
176184
| WrappingClause::Where
177185
| WrappingClause::Join { .. } => true,
178186

179-
WrappingClause::PolicyCheck => {
180-
ctx.before_cursor_matches_kind(&["="])
181-
&& matches!(f.kind, ProcKind::Function | ProcKind::Procedure)
187+
WrappingClause::CheckOrUsingClause => {
188+
!matches!(f.kind, ProcKind::Aggregate)
189+
&& (ctx.before_cursor_matches_kind(&["(", "keyword_and"])
190+
|| ctx.wrapping_node_kind.as_ref().is_some_and(|nk| {
191+
matches!(nk, WrappingNode::BinaryExpression)
192+
}))
182193
}
183194

184195
_ => false,
@@ -209,11 +220,21 @@ impl CompletionFilter<'_> {
209220
&& ctx.before_cursor_matches_kind(&["keyword_into"])
210221
}
211222

223+
WrappingClause::CreatePolicy
224+
| WrappingClause::AlterPolicy
225+
| WrappingClause::DropPolicy => {
226+
ctx.before_cursor_matches_kind(&["keyword_on"])
227+
}
228+
212229
_ => false,
213230
},
214231

215232
CompletionRelevanceData::Policy(_) => {
216-
matches!(clause, WrappingClause::PolicyName)
233+
matches!(
234+
clause,
235+
// not CREATE – there can't be existing policies.
236+
WrappingClause::AlterPolicy | WrappingClause::DropPolicy
237+
) && ctx.before_cursor_matches_kind(&["keyword_policy", "keyword_exists"])
217238
}
218239

219240
CompletionRelevanceData::Role(_) => match clause {
@@ -224,6 +245,11 @@ impl CompletionFilter<'_> {
224245
WrappingClause::SetStatement => ctx
225246
.before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
226247

248+
WrappingClause::AlterPolicy | WrappingClause::CreatePolicy => {
249+
ctx.before_cursor_matches_kind(&["keyword_to"])
250+
&& ctx.matches_ancestor_history(&["policy_to_role"])
251+
}
252+
227253
_ => false,
228254
},
229255
}

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

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ impl CompletionScore<'_> {
3535
self.check_matching_wrapping_node(ctx);
3636
self.check_relations_in_stmt(ctx);
3737
self.check_columns_in_stmt(ctx);
38+
self.check_is_not_wellknown_migration(ctx);
3839
}
3940

4041
fn check_matches_query_input(&mut self, ctx: &TreesitterContext) {
@@ -100,12 +101,14 @@ impl CompletionScore<'_> {
100101
WrappingClause::Select if !has_mentioned_tables => 15,
101102
WrappingClause::Select if has_mentioned_tables => 0,
102103
WrappingClause::From => 0,
104+
WrappingClause::CheckOrUsingClause => 0,
103105
_ => -50,
104106
},
105107
CompletionRelevanceData::Column(col) => match clause_type {
106108
WrappingClause::Select if has_mentioned_tables => 10,
107109
WrappingClause::Select if !has_mentioned_tables => 0,
108110
WrappingClause::Where => 10,
111+
WrappingClause::CheckOrUsingClause => 0,
109112
WrappingClause::Join { on_node }
110113
if on_node.is_some_and(|on| {
111114
ctx.node_under_cursor
@@ -123,10 +126,13 @@ impl CompletionScore<'_> {
123126
WrappingClause::Join { .. } if !has_mentioned_schema => 15,
124127
WrappingClause::Update if !has_mentioned_schema => 15,
125128
WrappingClause::Delete if !has_mentioned_schema => 15,
129+
WrappingClause::AlterPolicy if !has_mentioned_schema => 15,
130+
WrappingClause::DropPolicy if !has_mentioned_schema => 15,
131+
WrappingClause::CreatePolicy if !has_mentioned_schema => 15,
126132
_ => -50,
127133
},
128134
CompletionRelevanceData::Policy(_) => match clause_type {
129-
WrappingClause::PolicyName => 25,
135+
WrappingClause::AlterPolicy | WrappingClause::DropPolicy => 25,
130136
_ => -50,
131137
},
132138

@@ -156,6 +162,7 @@ impl CompletionScore<'_> {
156162
_ => -50,
157163
},
158164
CompletionRelevanceData::Function(_) => match wrapping_node {
165+
WrappingNode::BinaryExpression => 15,
159166
WrappingNode::Relation => 10,
160167
_ => -50,
161168
},
@@ -184,7 +191,6 @@ impl CompletionScore<'_> {
184191
}
185192

186193
fn check_matches_schema(&mut self, ctx: &TreesitterContext) {
187-
// TODO
188194
let schema_name = match ctx.schema_or_alias_name.as_ref() {
189195
None => return,
190196
Some(n) => n.replace('"', ""),
@@ -349,4 +355,18 @@ impl CompletionScore<'_> {
349355
}
350356
}
351357
}
358+
359+
fn check_is_not_wellknown_migration(&mut self, _ctx: &TreesitterContext) {
360+
if let Some(table_name) = self.get_table_name() {
361+
if ["_sqlx_migrations"].contains(&table_name) {
362+
self.score -= 10;
363+
}
364+
}
365+
366+
if let Some(schema_name) = self.get_schema_name() {
367+
if ["supabase_migrations"].contains(&schema_name) {
368+
self.score -= 10;
369+
}
370+
}
371+
}
352372
}

‎crates/pgt_hover/src/hovered_node.rs‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
use pgt_treesitter::WrappingClause;
2+
13
#[derive(Debug)]
24
pub(crate) enum NodeIdentification {
35
Name(String),
@@ -40,6 +42,28 @@ impl HoveredNode {
4042
Some(HoveredNode::Table(NodeIdentification::Name(node_content)))
4143
}
4244
}
45+
46+
"identifier"
47+
if ctx.matches_ancestor_history(&["object_reference"])
48+
&& ctx.wrapping_clause_type.as_ref().is_some_and(|clause| {
49+
matches!(
50+
clause,
51+
WrappingClause::AlterPolicy
52+
| WrappingClause::CreatePolicy
53+
| WrappingClause::DropPolicy
54+
)
55+
}) =>
56+
{
57+
if let Some(schema) = ctx.schema_or_alias_name.as_ref() {
58+
Some(HoveredNode::Table(NodeIdentification::SchemaAndName((
59+
schema.clone(),
60+
node_content,
61+
))))
62+
} else {
63+
Some(HoveredNode::Table(NodeIdentification::Name(node_content)))
64+
}
65+
}
66+
4367
"identifier" if ctx.matches_ancestor_history(&["field"]) => {
4468
if let Some(table_or_alias) = ctx.schema_or_alias_name.as_ref() {
4569
Some(HoveredNode::Column(NodeIdentification::SchemaAndName((
@@ -62,7 +86,7 @@ impl HoveredNode {
6286
)))
6387
}
6488
}
65-
"identifier" if ctx.matches_ancestor_history(&["alter_role"]) => {
89+
"identifier" if ctx.matches_one_of_ancestors(&["alter_role", "policy_to_role"]) => {
6690
Some(HoveredNode::Role(NodeIdentification::Name(node_content)))
6791
}
6892

@@ -75,7 +99,7 @@ impl HoveredNode {
7599
Some(HoveredNode::Column(NodeIdentification::Name(node_content)))
76100
}
77101

78-
"policy_table" | "revoke_table" | "grant_table" => {
102+
"revoke_table" | "grant_table" => {
79103
if let Some(schema) = ctx.schema_or_alias_name.as_ref() {
80104
Some(HoveredNode::Table(NodeIdentification::SchemaAndName((
81105
schema.clone(),

0 commit comments

Comments
 (0)