diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 363a1579..6924bf53 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/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 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); } diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 59e9c2e2..7fb537b6 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/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(" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 92a7f1ce..ce67925f 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -113,6 +113,7 @@ pub fn hover(file: &ast::SourceFile, offset: TextSize) -> Option { | NameRefClass::DeleteReturningQualifiedColumnTable | NameRefClass::MergeReturningQualifiedColumnTable | NameRefClass::MergeTable + | NameRefClass::PolicyQualifiedColumnTable | NameRefClass::ForeignKeyTable | NameRefClass::LikeTable | NameRefClass::InheritsTable diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 3400719c..f7379939 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/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()