ide: goto def with policy stmts

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;

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<NameRefClass>
         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<NameRefClass>
         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);
         }

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("

crates/squawk_ide/src/hover.rs

@@ -57,6 +57,7 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option<String> {
             | 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<String> {
             }
             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<String> {
+    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<String> {
+    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<String> {
     let name = create_event_trigger.name()?.syntax().text().to_string();
     Some(format!("event trigger {}", name))

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<Schema>,
+    position: TextSize,
+    table: Name,
+) -> Option<SyntaxNodePtr> {
+    binder.lookup_with_table(
+        policy_name,
+        SymbolKind::Policy,
+        position,
+        schema,
+        &Some(table),
+    )
+}
+
 fn resolve_event_trigger_name_ptr(
     binder: &Binder,
     event_trigger_name: &Name,

crates/squawk_ide/src/symbols.rs

@@ -63,6 +63,7 @@ pub(crate) enum SymbolKind {
     Trigger,
     EventTrigger,
     Role,
+    Policy,
 }
 
 #[derive(Clone, Debug)]