diff --git a/crates/squawk_ide/src/binder.rs b/crates/squawk_ide/src/binder.rs index 271813e3..4ee11834 100644 --- a/crates/squawk_ide/src/binder.rs +++ b/crates/squawk_ide/src/binder.rs @@ -279,6 +279,7 @@ fn bind_stmt(b: &mut Binder, stmt: ast::Stmt) { ast::Stmt::Prepare(prepare) => bind_prepare(b, prepare), ast::Stmt::Listen(listen) => bind_listen(b, listen), ast::Stmt::Set(set) => bind_set(b, set), + ast::Stmt::CreatePolicy(create_policy) => bind_create_policy(b, create_policy), _ => {} } } @@ -696,6 +697,38 @@ fn bind_create_trigger(b: &mut Binder, create_trigger: ast::CreateTrigger) { b.scopes[root].insert(trigger_name, trigger_id); } +fn bind_create_policy(b: &mut Binder, create_policy: ast::CreatePolicy) { + let Some(name) = create_policy.name() else { + return; + }; + + let policy_name = Name::from_node(&name); + let name_ptr = SyntaxNodePtr::new(name.syntax()); + + let Some(table_path) = create_policy.on_table().and_then(|on| on.path()) else { + return; + }; + + let Some(table_name) = item_name(&table_path) else { + return; + }; + + let Some(schema) = schema_name(b, &table_path, false) else { + return; + }; + + let policy_id = b.symbols.alloc(Symbol { + kind: SymbolKind::Policy, + ptr: name_ptr, + schema: Some(schema), + params: None, + table: Some(table_name), + }); + + let root = b.root_scope(); + b.scopes[root].insert(policy_name, policy_id); +} + fn bind_create_event_trigger(b: &mut Binder, create_event_trigger: ast::CreateEventTrigger) { let Some(name) = create_event_trigger.name() else { return; diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 6c856b31..363a1579 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/crates/squawk_ide/src/classify.rs @@ -16,6 +16,8 @@ pub(crate) enum NameRefClass { DropMaterializedView, DropSequence, DropTrigger, + DropPolicy, + AlterPolicy, DropEventTrigger, SequenceOwnedByColumn, Tablespace, @@ -34,6 +36,7 @@ pub(crate) enum NameRefClass { ForeignKeyColumn, ForeignKeyLocalColumn, CheckConstraintColumn, + PolicyColumn, GeneratedColumn, UniqueConstraintColumn, PrimaryKeyConstraintColumn, @@ -474,6 +477,12 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option if ast::DropTrigger::can_cast(ancestor.kind()) { return Some(NameRefClass::DropTrigger); } + if ast::DropPolicy::can_cast(ancestor.kind()) { + return Some(NameRefClass::DropPolicy); + } + if ast::AlterPolicy::can_cast(ancestor.kind()) { + return Some(NameRefClass::AlterPolicy); + } if ast::DropEventTrigger::can_cast(ancestor.kind()) { return Some(NameRefClass::DropEventTrigger); } @@ -579,6 +588,11 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option if ast::CheckConstraint::can_cast(ancestor.kind()) { return Some(NameRefClass::CheckConstraintColumn); } + if ast::CreatePolicy::can_cast(ancestor.kind()) + || ast::AlterPolicy::can_cast(ancestor.kind()) + { + return Some(NameRefClass::PolicyColumn); + } if ast::GeneratedConstraint::can_cast(ancestor.kind()) { return Some(NameRefClass::GeneratedColumn); } diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 71272f1c..59e9c2e2 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -500,6 +500,72 @@ drop trigger tr$0 on t; "); } + #[test] + fn goto_drop_policy() { + assert_snapshot!(goto(" +create table t(c int); +create table u(c int); +create policy p on t; +create policy p on u; +drop policy if exists p$0 on t; +"), @r" + ╭▸ + 4 │ create policy p on t; + │ ─ 2. destination + 5 │ create policy p on u; + 6 │ drop policy if exists p on t; + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_alter_policy() { + assert_snapshot!(goto(" +create table t(c int); +create policy p on t; +alter policy p$0 on t + with check (c > 1); +"), @r" + ╭▸ + 3 │ create policy p on t; + │ ─ 2. destination + 4 │ alter policy p on t + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_create_policy_column() { + assert_snapshot!(goto(" +create table t(c int, d int); +create policy p on t + with check (c$0 > d); +"), @r" + ╭▸ + 2 │ create table t(c int, d int); + │ ─ 2. destination + 3 │ create policy p on t + 4 │ with check (c > d); + ╰╴ ─ 1. source + "); + } + + #[test] + fn goto_create_policy_using_column() { + assert_snapshot!(goto(" +create table t(c int, d int); +create policy p on t + using (c$0 > d and 1 < 2); +"), @r" + ╭▸ + 2 │ create table t(c int, d int); + │ ─ 2. destination + 3 │ create policy p on t + 4 │ using (c > d and 1 < 2); + ╰╴ ─ 1. source + "); + } + #[test] fn goto_drop_event_trigger() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 0bbdb04d..92a7f1ce 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -57,6 +57,7 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option { | NameRefClass::MergeWhenColumn | NameRefClass::MergeOnColumn | NameRefClass::CheckConstraintColumn + | NameRefClass::PolicyColumn | NameRefClass::GeneratedColumn | NameRefClass::UniqueConstraintColumn | NameRefClass::PrimaryKeyConstraintColumn @@ -127,6 +128,9 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option { } NameRefClass::DropSequence => return hover_sequence(root, &name_ref, &binder), NameRefClass::DropTrigger => return hover_trigger(root, &name_ref, &binder), + NameRefClass::DropPolicy | NameRefClass::AlterPolicy => { + return hover_policy(root, &name_ref, &binder); + } NameRefClass::DropEventTrigger | NameRefClass::AlterEventTrigger => { return hover_event_trigger(root, &name_ref, &binder); } @@ -815,6 +819,24 @@ fn hover_trigger( format_create_trigger(&create_trigger, binder) } +fn hover_policy( + root: &SyntaxNode, + name_ref: &ast::NameRef, + binder: &binder::Binder, +) -> Option { + let policy_ptr = resolve::resolve_name_ref_ptrs(binder, root, name_ref)? + .into_iter() + .next()?; + + let policy_name_node = policy_ptr.to_node(root); + + let create_policy = policy_name_node + .ancestors() + .find_map(ast::CreatePolicy::cast)?; + + format_create_policy(&create_policy, binder) +} + fn hover_event_trigger( root: &SyntaxNode, name_ref: &ast::NameRef, @@ -1103,6 +1125,20 @@ fn format_create_trigger( )) } +fn format_create_policy( + create_policy: &ast::CreatePolicy, + binder: &binder::Binder, +) -> Option { + let policy_name = create_policy.name()?.syntax().text().to_string(); + let on_table_path = create_policy.on_table()?.path()?; + + let (schema, table_name) = resolve::resolve_table_info(binder, &on_table_path)?; + Some(format!( + "policy {}.{} on {}.{}", + schema, policy_name, schema, table_name + )) +} + fn format_create_event_trigger(create_event_trigger: &ast::CreateEventTrigger) -> Option { let name = create_event_trigger.name()?.syntax().text().to_string(); Some(format!("event trigger {}", name)) diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index c13ba63f..3400719c 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -173,6 +173,22 @@ pub(crate) fn resolve_name_ref_ptrs( resolve_trigger_name_ptr(binder, &trigger_name, &schema, position, Some(table_name)) .map(|ptr| smallvec![ptr]) } + NameRefClass::DropPolicy | NameRefClass::AlterPolicy => { + let (policy_name, on_table) = name_ref.syntax().ancestors().find_map(|a| { + if let Some(policy) = ast::DropPolicy::cast(a.clone()) { + Some((policy.name_ref(), policy.on_table())) + } else { + ast::AlterPolicy::cast(a).map(|policy| (policy.name_ref(), policy.on_table())) + } + })?; + let policy_name = Name::from_node(&policy_name?); + let on_table_path = on_table.and_then(|on_table| on_table.path())?; + let schema = extract_schema_name(&on_table_path); + let table_name = extract_table_name(&on_table_path)?; + let position = name_ref.syntax().text_range().start(); + resolve_policy_name_ptr(binder, &policy_name, &schema, position, table_name) + .map(|ptr| smallvec![ptr]) + } NameRefClass::DropEventTrigger | NameRefClass::AlterEventTrigger => { let event_trigger_name = Name::from_node(name_ref); resolve_event_trigger_name_ptr(binder, &event_trigger_name).map(|ptr| smallvec![ptr]) @@ -259,6 +275,20 @@ pub(crate) fn resolve_name_ref_ptrs( find_column_in_create_table(binder, root, &create_table, &column_name) .map(|ptr| smallvec![ptr]) } + NameRefClass::PolicyColumn => { + let on_table_path = name_ref.syntax().ancestors().find_map(|n| { + if let Some(create_policy) = ast::CreatePolicy::cast(n.clone()) { + create_policy.on_table()?.path() + } else if let Some(alter_policy) = ast::AlterPolicy::cast(n) { + alter_policy.on_table()?.path() + } else { + None + } + })?; + let column_name = Name::from_node(name_ref); + resolve_column_for_path(binder, root, &on_table_path, column_name) + .map(|ptr| smallvec![ptr]) + } NameRefClass::LikeTable => { let like_clause = name_ref .syntax() @@ -653,6 +683,22 @@ fn resolve_trigger_name_ptr( binder.lookup_with_table(trigger_name, SymbolKind::Trigger, position, schema, &table) } +fn resolve_policy_name_ptr( + binder: &Binder, + policy_name: &Name, + schema: &Option, + position: TextSize, + table: Name, +) -> Option { + binder.lookup_with_table( + policy_name, + SymbolKind::Policy, + position, + schema, + &Some(table), + ) +} + fn resolve_event_trigger_name_ptr( binder: &Binder, event_trigger_name: &Name, diff --git a/crates/squawk_ide/src/symbols.rs b/crates/squawk_ide/src/symbols.rs index 7d8b3b3f..b32885d0 100644 --- a/crates/squawk_ide/src/symbols.rs +++ b/crates/squawk_ide/src/symbols.rs @@ -63,6 +63,7 @@ pub(crate) enum SymbolKind { Trigger, EventTrigger, Role, + Policy, } #[derive(Clone, Debug)]