ide: goto def on columns in policy stmts

crates/squawk_ide/src/classify.rs

@@ -37,6 +37,7 @@ pub(crate) enum NameRefClass {
     ForeignKeyLocalColumn,
     CheckConstraintColumn,
     PolicyColumn,
+    PolicyQualifiedColumnTable,
     GeneratedColumn,
     UniqueConstraintColumn,
     PrimaryKeyConstraintColumn,
@@ -275,6 +276,15 @@ pub(crate) fn classify_name_ref(name_ref: &ast::NameRef) -> Option<NameRefClass>
                     return Some(NameRefClass::SelectQualifiedColumnTable);
                 }
             }
+            if ast::CreatePolicy::can_cast(ancestor.kind())
+                || ast::AlterPolicy::can_cast(ancestor.kind())
+            {
+                if is_function_call || is_schema_table_col {
+                    return Some(NameRefClass::SchemaQualifier);
+                } else {
+                    return Some(NameRefClass::PolicyQualifiedColumnTable);
+                }
+            }
         }
         return Some(NameRefClass::SchemaQualifier);
     }

crates/squawk_ide/src/goto_definition.rs

@@ -534,6 +534,23 @@ alter policy p$0 on t
         ");
     }
 
+    #[test]
+    fn goto_alter_policy_column() {
+        assert_snapshot!(goto("
+create table t(c int);
+create policy p on t;
+alter policy p on t
+  with check (c$0 > 1);
+"), @r"
+          ╭▸ 
+        3 │ create policy p on t;
+          │               ─ 2. destination
+        4 │ alter policy p on t
+        5 │   with check (c > 1);
+          ╰╴              ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_create_policy_column() {
         assert_snapshot!(goto("
@@ -566,6 +583,54 @@ create policy p on t
         ");
     }
 
+    #[test]
+    fn goto_create_policy_qualified_column_table() {
+        assert_snapshot!(goto("
+create table t(c int, d int);
+create policy p on t
+  with check (t$0.c > d);
+"), @r"
+          ╭▸ 
+        2 │ create table t(c int, d int);
+          │              ─ 2. destination
+        3 │ create policy p on t
+        4 │   with check (t.c > d);
+          ╰╴              ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_create_policy_qualified_column() {
+        assert_snapshot!(goto("
+create table t(c int, d int);
+create policy p on t
+  with check (t.c$0 > d);
+"), @r"
+          ╭▸ 
+        2 │ create table t(c int, d int);
+          │                ─ 2. destination
+        3 │ create policy p on t
+        4 │   with check (t.c > d);
+          ╰╴                ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_alter_policy_qualified_column_table() {
+        assert_snapshot!(goto("
+create table t(c int, d int);
+alter policy p on t
+  with check (t$0.c > d);
+"), @r"
+          ╭▸ 
+        2 │ create table t(c int, d int);
+          │              ─ 2. destination
+        3 │ alter policy p on t
+        4 │   with check (t.c > d);
+          ╰╴              ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_drop_event_trigger() {
         assert_snapshot!(goto("

crates/squawk_ide/src/hover.rs

@@ -113,6 +113,7 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option<String> {
             | NameRefClass::DeleteReturningQualifiedColumnTable
             | NameRefClass::MergeReturningQualifiedColumnTable
             | NameRefClass::MergeTable
+            | NameRefClass::PolicyQualifiedColumnTable
             | NameRefClass::ForeignKeyTable
             | NameRefClass::LikeTable
             | NameRefClass::InheritsTable

crates/squawk_ide/src/resolve.rs

@@ -289,6 +289,21 @@ pub(crate) fn resolve_name_ref_ptrs(
             resolve_column_for_path(binder, root, &on_table_path, column_name)
                 .map(|ptr| smallvec![ptr])
         }
+        NameRefClass::PolicyQualifiedColumnTable => {
+            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 table_name = extract_table_name(&on_table_path)?;
+            let schema = extract_schema_name(&on_table_path);
+            let position = name_ref.syntax().text_range().start();
+            resolve_table_name_ptr(binder, &table_name, &schema, position).map(|ptr| smallvec![ptr])
+        }
         NameRefClass::LikeTable => {
             let like_clause = name_ref
                 .syntax()