Skip to content

Commit c0ed76c

Browse files
juleswritescodejuleswritescode
committed
ok?
1 parent 4a62cc0 commit c0ed76c

File tree

10 files changed

+316392
-313222
lines changed

10 files changed

+316392
-313222
lines changed

‎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/relevance/filtering.rs‎

Lines changed: 21 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,
@@ -162,8 +160,11 @@ impl CompletionFilter<'_> {
162160
&& ctx.matches_ancestor_history(&["field"]))
163161
}
164162

165-
WrappingClause::PolicyCheck => {
166-
ctx.before_cursor_matches_kind(&["keyword_and", "("])
163+
WrappingClause::CheckOrUsingClause => {
164+
ctx.before_cursor_matches_kind(&["(", "keyword_and"])
165+
|| ctx.wrapping_node_kind.as_ref().is_some_and(|nk| {
166+
matches!(nk, WrappingNode::BinaryExpression)
167+
})
167168
}
168169

169170
_ => false,
@@ -176,9 +177,12 @@ impl CompletionFilter<'_> {
176177
| WrappingClause::Where
177178
| WrappingClause::Join { .. } => true,
178179

179-
WrappingClause::PolicyCheck => {
180-
ctx.before_cursor_matches_kind(&["="])
181-
&& matches!(f.kind, ProcKind::Function | ProcKind::Procedure)
180+
WrappingClause::CheckOrUsingClause => {
181+
!matches!(f.kind, ProcKind::Aggregate)
182+
&& (ctx.before_cursor_matches_kind(&["(", "keyword_and"])
183+
|| ctx.wrapping_node_kind.as_ref().is_some_and(|nk| {
184+
matches!(nk, WrappingNode::BinaryExpression)
185+
}))
182186
}
183187

184188
_ => false,
@@ -213,7 +217,11 @@ impl CompletionFilter<'_> {
213217
},
214218

215219
CompletionRelevanceData::Policy(_) => {
216-
matches!(clause, WrappingClause::PolicyName)
220+
matches!(
221+
clause,
222+
// not CREATE – there can't be existing policies.
223+
WrappingClause::DropPolicy | WrappingClause::AlterPolicy
224+
)
217225
}
218226

219227
CompletionRelevanceData::Role(_) => match clause {
@@ -224,6 +232,10 @@ impl CompletionFilter<'_> {
224232
WrappingClause::SetStatement => ctx
225233
.before_cursor_matches_kind(&["keyword_role", "keyword_authorization"]),
226234

235+
WrappingClause::AlterPolicy | WrappingClause::CreatePolicy => {
236+
ctx.before_cursor_matches_kind(&["keyword_to"])
237+
}
238+
227239
_ => false,
228240
},
229241
}

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

Lines changed: 19 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
@@ -126,7 +129,7 @@ impl CompletionScore<'_> {
126129
_ => -50,
127130
},
128131
CompletionRelevanceData::Policy(_) => match clause_type {
129-
WrappingClause::PolicyName => 25,
132+
WrappingClause::AlterPolicy | WrappingClause::DropPolicy => 25,
130133
_ => -50,
131134
},
132135

@@ -156,6 +159,7 @@ impl CompletionScore<'_> {
156159
_ => -50,
157160
},
158161
CompletionRelevanceData::Function(_) => match wrapping_node {
162+
WrappingNode::BinaryExpression => 15,
159163
WrappingNode::Relation => 10,
160164
_ => -50,
161165
},
@@ -184,7 +188,6 @@ impl CompletionScore<'_> {
184188
}
185189

186190
fn check_matches_schema(&mut self, ctx: &TreesitterContext) {
187-
// TODO
188191
let schema_name = match ctx.schema_or_alias_name.as_ref() {
189192
None => return,
190193
Some(n) => n.replace('"', ""),
@@ -349,4 +352,18 @@ impl CompletionScore<'_> {
349352
}
350353
}
351354
}
355+
356+
fn check_is_not_wellknown_migration(&mut self, _ctx: &TreesitterContext) {
357+
if let Some(table_name) = self.get_table_name() {
358+
if ["_sqlx_migrations"].contains(&table_name) {
359+
self.score -= 10;
360+
}
361+
}
362+
363+
if let Some(schema_name) = self.get_schema_name() {
364+
if ["supabase_migrations"].contains(&schema_name) {
365+
self.score -= 10;
366+
}
367+
}
368+
}
352369
}

‎crates/pgt_hover/src/hovered_node.rs‎

Lines changed: 25 additions & 1 deletion
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((
@@ -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(),

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

Lines changed: 17 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,13 @@ use std::{
44
};
55
mod base_parser;
66
mod grant_parser;
7-
mod policy_parser;
87
mod revoke_parser;
98

109
use crate::queries::{self, QueryResult, TreeSitterQueriesExecutor};
1110
use pgt_text_size::{TextRange, TextSize};
1211

1312
use crate::context::{
14-
base_parser::CompletionStatementParser,
15-
grant_parser::GrantParser,
16-
policy_parser::{PolicyParser, PolicyStmtKind},
17-
revoke_parser::RevokeParser,
13+
base_parser::CompletionStatementParser, grant_parser::GrantParser, revoke_parser::RevokeParser,
1814
};
1915

2016
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
@@ -34,20 +30,15 @@ pub enum WrappingClause<'a> {
3430
DropColumn,
3531
AlterColumn,
3632
RenameColumn,
37-
PolicyName,
3833
ToRoleAssignment,
3934
SetStatement,
4035
AlterRole,
4136
DropRole,
4237

43-
/// `PolicyCheck` refers to either the `WITH CHECK` or the `USING` clause
44-
/// in a policy statement.
45-
/// ```sql
46-
/// CREATE POLICY "my pol" ON PUBLIC.USERS
47-
/// FOR SELECT
48-
/// USING (...) -- this one!
49-
/// ```
50-
PolicyCheck,
38+
CreatePolicy,
39+
AlterPolicy,
40+
DropPolicy,
41+
CheckOrUsingClause,
5142
}
5243

5344
#[derive(PartialEq, Eq, Hash, Debug, Clone)]
@@ -205,12 +196,7 @@ impl<'a> TreesitterContext<'a> {
205196
mentioned_columns: HashMap::new(),
206197
};
207198

208-
// policy handling is important to Supabase, but they are a PostgreSQL specific extension,
209-
// so the pgt_treesitter_grammar language does not support it.
210-
// We infer the context manually.
211-
if PolicyParser::looks_like_matching_stmt(params.text) {
212-
ctx.gather_policy_context();
213-
} else if GrantParser::looks_like_matching_stmt(params.text) {
199+
if GrantParser::looks_like_matching_stmt(params.text) {
214200
ctx.gather_grant_context();
215201
} else if RevokeParser::looks_like_matching_stmt(params.text) {
216202
ctx.gather_revoke_context();
@@ -278,43 +264,6 @@ impl<'a> TreesitterContext<'a> {
278264
};
279265
}
280266

281-
fn gather_policy_context(&mut self) {
282-
let policy_context = PolicyParser::get_context(self.text, self.position);
283-
284-
self.node_under_cursor = Some(NodeUnderCursor::CustomNode {
285-
text: policy_context.node_text,
286-
range: policy_context.node_range,
287-
kind: policy_context.node_kind.clone(),
288-
previous_node_kind: Some(policy_context.previous_node_kind),
289-
});
290-
291-
if policy_context.node_kind == "policy_table" {
292-
self.schema_or_alias_name = policy_context.schema_name.clone();
293-
}
294-
295-
if policy_context.table_name.is_some() {
296-
let mut new = HashSet::new();
297-
new.insert(policy_context.table_name.unwrap());
298-
self.mentioned_relations
299-
.insert(policy_context.schema_name, new);
300-
}
301-
302-
self.wrapping_clause_type = match policy_context.node_kind.as_str() {
303-
"policy_name" if policy_context.statement_kind != PolicyStmtKind::Create => {
304-
Some(WrappingClause::PolicyName)
305-
}
306-
"policy_role" => Some(WrappingClause::ToRoleAssignment),
307-
"policy_table" => Some(WrappingClause::From),
308-
_ => {
309-
if policy_context.in_check_or_using_clause {
310-
Some(WrappingClause::PolicyCheck)
311-
} else {
312-
None
313-
}
314-
}
315-
};
316-
}
317-
318267
fn gather_info_from_ts_queries(&mut self) {
319268
let stmt_range = self.wrapping_statement_range.as_ref();
320269
let sql = self.text;
@@ -498,13 +447,6 @@ impl<'a> TreesitterContext<'a> {
498447
}
499448
}
500449

501-
"where" | "update" | "select" | "delete" | "from" | "join" | "column_definitions"
502-
| "alter_role" | "drop_role" | "set_statement" | "drop_table" | "alter_table"
503-
| "drop_column" | "alter_column" | "rename_column" => {
504-
self.wrapping_clause_type =
505-
self.get_wrapping_clause_from_current_node(current_node, &mut cursor);
506-
}
507-
508450
"relation" | "binary_expression" | "assignment" => {
509451
self.wrapping_node_kind = current_node_kind.try_into().ok();
510452
}
@@ -518,7 +460,13 @@ impl<'a> TreesitterContext<'a> {
518460
}
519461
}
520462

521-
_ => {}
463+
_ => {
464+
if let Some(clause_type) =
465+
self.get_wrapping_clause_from_current_node(current_node, &mut cursor)
466+
{
467+
self.wrapping_clause_type = Some(clause_type);
468+
}
469+
}
522470
}
523471

524472
// We have arrived at the leaf node
@@ -734,6 +682,10 @@ impl<'a> TreesitterContext<'a> {
734682
"alter_table" => Some(WrappingClause::AlterTable),
735683
"set_statement" => Some(WrappingClause::SetStatement),
736684
"column_definitions" => Some(WrappingClause::ColumnDefinitions),
685+
"create_policy" => Some(WrappingClause::CreatePolicy),
686+
"alter_policy" => Some(WrappingClause::AlterPolicy),
687+
"drop_policy" => Some(WrappingClause::DropPolicy),
688+
"check_or_using_clause" => Some(WrappingClause::CheckOrUsingClause),
737689
"insert" => Some(WrappingClause::Insert),
738690
"join" => {
739691
// sadly, we need to manually iterate over the children –

0 commit comments

Comments
 (0)