ide: refactor resolve and fix a couple edge cases

crates/squawk_ide/src/classify.rs

@@ -28,6 +28,7 @@ pub(crate) enum NameRefClass {
     Index,
     InsertColumn,
     InsertQualifiedColumnTable,
+    InsertTable,
     JoinUsingColumn,
     LikeTable,
     MergeColumn,
@@ -706,7 +707,7 @@ pub(crate) fn classify_name_ref(node: &SyntaxNode) -> Option<NameRefClass> {
             {
                 return Some(NameRefClass::InsertColumn);
             }
-            return Some(NameRefClass::Table);
+            return Some(NameRefClass::InsertTable);
         }
         if ast::JoinUsingClause::can_cast(ancestor.kind()) && in_column_list {
             return Some(NameRefClass::JoinUsingColumn);

crates/squawk_ide/src/collect.rs

@@ -1,14 +1,14 @@
 use crate::ast_nav;
 use crate::builtins::builtins_file;
 use crate::column_name::ColumnName;
-use crate::db::{File, bind, parse};
+use crate::db::{File, parse};
 use crate::infer::{Type, infer_type_from_expr, infer_type_from_ty};
+use crate::location::LocationKind;
 use crate::name::{self, Name};
 use crate::resolve::{
-    ResolvedTableName, find_from_item_in_from_clause, qualified_star_table_name, resolve_cte_table,
-    resolve_table_name, table_ptr_from_from_item,
+    ResolvedTableName, find_from_item_in_from_clause, qualified_star_table_name,
+    resolve_table_like, resolve_table_name, table_ptr_from_from_item,
 };
-use crate::symbols::SymbolKind;
 use salsa::Database as Db;
 use squawk_syntax::{
     SyntaxNodePtr,
@@ -227,11 +227,7 @@ pub(crate) fn view_like_columns_with_types(
     results
 }
 
-pub(crate) fn with_table_column_names(
-    db: &dyn Db,
-    file: File,
-    with_table: ast::WithTable,
-) -> Vec<Name> {
+fn with_table_column_names(db: &dyn Db, file: File, with_table: ast::WithTable) -> Vec<Name> {
     with_table_columns_with_types(db, file, with_table)
         .into_iter()
         .map(|(name, _)| name)
@@ -435,9 +431,7 @@ fn columns_for_star_from_table_ptr(
     }
 }
 
-pub(crate) fn target_list_columns_with_types(
-    target_list: &ast::TargetList,
-) -> Vec<(Name, Option<Type>)> {
+fn target_list_columns_with_types(target_list: &ast::TargetList) -> Vec<(Name, Option<Type>)> {
     let mut columns = vec![];
     for target in target_list.targets() {
         if let Some((col_name, _node)) = ColumnName::from_target(target.clone())
@@ -466,7 +460,6 @@ fn select_variant_columns_with_types(
     file: File,
     select_variant: &ast::SelectVariant,
 ) -> Vec<(Name, Option<Type>)> {
-    let binder = bind(db, file);
     let tree = parse(db, file).tree();
     let root = tree.syntax();
     match select_variant {
@@ -517,35 +510,34 @@ fn select_variant_columns_with_types(
             let Some((schema, table_name)) = name::schema_and_name_path(&path) else {
                 return vec![];
             };
-
-            if schema.is_none()
-                && let Some(name_ref) = path.segment().and_then(|segment| segment.name_ref())
-                && let Some(with_table_ptr) = resolve_cte_table(&name_ref, &table_name)
-                && let Some(with_table) = with_table_ptr
-                    .to_node(root)
-                    .ancestors()
-                    .find_map(ast::WithTable::cast)
-            {
-                return with_table_column_names(db, file, with_table)
-                    .into_iter()
-                    .map(|name| (name, None))
-                    .collect();
-            }
-
+            let name_ref = path.segment().and_then(|segment| segment.name_ref());
             let position = table.syntax().text_range().start();
-            let Some(table_ptr) =
-                binder.lookup_with(&table_name, SymbolKind::Table, position, &schema)
+            let Some((ptr, kind)) =
+                resolve_table_like(db, file, name_ref.as_ref(), &table_name, &schema, position)
             else {
                 return vec![];
             };
-            let Some(create_table) = table_ptr
-                .to_node(root)
-                .ancestors()
-                .find_map(ast::CreateTableLike::cast)
-            else {
-                return vec![];
-            };
-            table_columns(db, file, &create_table)
+            let node = ptr.to_node(root);
+            match kind {
+                LocationKind::View => node
+                    .ancestors()
+                    .find_map(ast::CreateViewLike::cast)
+                    .map(|v| view_like_columns_with_types(&v))
+                    .unwrap_or_default(),
+                LocationKind::Table => {
+                    if let Some(with_table) = node.ancestors().find_map(ast::WithTable::cast) {
+                        return with_table_column_names(db, file, with_table)
+                            .into_iter()
+                            .map(|name| (name, None))
+                            .collect();
+                    }
+                    node.ancestors()
+                        .find_map(ast::CreateTableLike::cast)
+                        .map(|t| table_columns(db, file, &t))
+                        .unwrap_or_default()
+                }
+                _ => vec![],
+            }
         }
     }
 }

crates/squawk_ide/src/goto_definition.rs

@@ -533,6 +533,24 @@ select t$0.* from t;
         ");
     }
 
+    #[test]
+    fn goto_cte_shadowed_table_star_column_count_not_found() {
+        goto_not_found(
+            "
+create table t(a int, b int);
+with
+  t as (
+    select 1
+  ),
+  -- yy overrides y since there's only 1 column in the *
+  u(x, yy) as (
+    select *, 2 y, 3 z from t
+  )
+select y$0 from u;
+",
+        );
+    }
+
     #[test]
     fn goto_cross_join_func_column() {
         assert_snapshot!(goto(r#"
@@ -4192,6 +4210,34 @@ select f2$0 from dup(42) as u(x, y);
         );
     }
 
+    #[test]
+    fn goto_fn_call_column_from_cte() {
+        assert_snapshot!(goto("
+with cte as (select 1 as a)
+select a$0(cte) from cte;
+"), @"
+          ╭▸ 
+        2 │ with cte as (select 1 as a)
+          │                          ─ 2. destination
+        3 │ select a(cte) from cte;
+          ╰╴       ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_fn_call_column_from_view() {
+        assert_snapshot!(goto("
+create view v as select 1 as a;
+select a$0(v) from v;
+"), @"
+          ╭▸ 
+        2 │ create view v as select 1 as a;
+          │                              ─ 2. destination
+        3 │ select a(v) from v;
+          ╰╴       ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_select_aggregate_call() {
         assert_snapshot!(goto("
@@ -5341,6 +5387,26 @@ insert into public.users$0(id, email) values (1, 'test@example.com');
         ");
     }
 
+    #[test]
+    fn goto_insert_view() {
+        assert_snapshot!(goto("
+create table users as select 1 id, 'joe' name, 'joe@example.com' email, 'active' status;
+create view active_users as
+  select id, name, email
+  from users
+  where status = 'active';
+insert into active_users$0 (name, email)
+values ('Alice', 'alice@example.com');
+"), @"
+          ╭▸ 
+        3 │ create view active_users as
+          │             ──────────── 2. destination
+          ‡
+        7 │ insert into active_users (name, email)
+          ╰╴                       ─ 1. source
+        ");
+    }
+
     #[test]
     fn goto_insert_column() {
         assert_snapshot!(goto("

crates/squawk_ide/src/hover.rs

@@ -5,12 +5,14 @@ use crate::column_name::ColumnName;
 use crate::comments::preceding_comment;
 use crate::db::{File, bind, parse};
 use crate::location::{Location, LocationKind};
+use crate::name;
 use crate::offsets::token_from_offset;
 use crate::symbols::{Name, Schema};
 use crate::{goto_definition, resolve};
 use rowan::TextSize;
 use salsa::Database as Db;
 use squawk_syntax::SyntaxNode;
+use squawk_syntax::SyntaxNodePtr;
 use squawk_syntax::{
     SyntaxKind,
     ast::{self, AstNode},
@@ -467,7 +469,7 @@ fn format_alias_with_column_list(db: &dyn Db, file: File, alias: ast::Alias) ->
 }
 
 fn hover_qualified_star(db: &dyn Db, file: File, field_expr: ast::FieldExpr) -> Option<Hover> {
-    let table_ptr = resolve::resolve_qualified_star_table_ptr(db, file, field_expr)?;
+    let table_ptr = qualified_star_table_ptr(db, file, field_expr)?;
     hover_qualified_star_columns(db, file, &table_ptr)
 }
 
@@ -484,7 +486,7 @@ fn hover_unqualified_star(db: &dyn Db, file: File, target: ast::Target) -> Optio
 fn hover_unqualified_star_with_binder(db: &dyn Db, file: File, target: &ast::Target) -> Vec<Hover> {
     let mut results = vec![];
 
-    if let Some(table_ptrs) = resolve::resolve_unqualified_star_table_ptrs(db, file, target) {
+    if let Some(table_ptrs) = unqualified_star_table_ptrs(db, file, target) {
         for table_ptr in table_ptrs {
             if let Some(columns) = hover_qualified_star_columns(db, file, &table_ptr) {
                 results.push(columns);
@@ -517,7 +519,7 @@ fn hover_unqualified_star_in_arg_list(
     file: File,
     arg_list: ast::ArgList,
 ) -> Option<Hover> {
-    let table_ptrs = resolve::resolve_unqualified_star_in_arg_list_ptrs(db, file, &arg_list)?;
+    let table_ptrs = unqualified_star_in_arg_list_ptrs(db, file, &arg_list)?;
     let mut results = vec![];
     for table_ptr in table_ptrs {
         if let Some(columns) = hover_qualified_star_columns(db, file, &table_ptr) {
@@ -736,7 +738,7 @@ fn hover_qualified_star_columns_from_subquery(
 
         for target in target_list.targets() {
             if target.star_token().is_some() {
-                let table_ptrs = resolve::resolve_unqualified_star_table_ptrs(db, file, &target)?;
+                let table_ptrs = unqualified_star_table_ptrs(db, file, &target)?;
                 for table_ptr in table_ptrs {
                     if let Some(columns) = hover_qualified_star_columns(db, file, &table_ptr) {
                         results.push(columns)
@@ -1394,6 +1396,123 @@ fn hover_routine(db: &dyn Db, def: Location) -> Option<Hover> {
     None
 }
 
+fn qualified_star_table_ptr(
+    db: &dyn Db,
+    file: File,
+    field_expr: ast::FieldExpr,
+) -> Option<SyntaxNodePtr> {
+    let table_name = resolve::qualified_star_table_name(&field_expr)?;
+    let position = field_expr.syntax().text_range().start();
+    let target = field_expr
+        .syntax()
+        .ancestors()
+        .find_map(ast::Target::cast)?;
+
+    let path = match ast_nav::target_parent_query(target)? {
+        ast_nav::ParentQuery::Select(select) => {
+            let from_clause = select.from_clause()?;
+            let from_item = resolve::find_from_item_in_from_clause(&from_clause, &table_name)?;
+
+            if let Some(alias) = from_item.alias()
+                && alias.column_list().is_some()
+            {
+                return Some(SyntaxNodePtr::new(alias.syntax()));
+            }
+
+            let (schema, table_name) = name::schema_and_table_from_from_item(&from_item)?;
+
+            let name_ref = from_item.name_ref();
+            if let Some((table_like_ptr, _kind)) = resolve::resolve_table_like(
+                db,
+                file,
+                name_ref.as_ref(),
+                &table_name,
+                &schema,
+                position,
+            ) {
+                return Some(table_like_ptr);
+            }
+
+            return None;
+        }
+        ast_nav::ParentQuery::Update(update) => update.relation_name()?.path()?,
+        ast_nav::ParentQuery::Delete(delete) => delete.relation_name()?.path()?,
+        ast_nav::ParentQuery::Insert(insert) => insert.path()?,
+        ast_nav::ParentQuery::Merge(merge) => merge.relation_name()?.path()?,
+    };
+
+    table_or_view_or_cte_ptrs(db, file, position, &path)?
+        .into_iter()
+        .next()
+}
+
+fn table_or_view_or_cte_ptrs(
+    db: &dyn Db,
+    file: File,
+    position: TextSize,
+    path: &ast::Path,
+) -> Option<Vec<SyntaxNodePtr>> {
+    let (schema, table_name) = name::schema_and_name_path(path)?;
+    let mut results = vec![];
+    let name_ref = path.segment().and_then(|x| x.name_ref());
+
+    if let Some((table_like_ptr, _kind)) =
+        resolve::resolve_table_like(db, file, name_ref.as_ref(), &table_name, &schema, position)
+    {
+        results.push(table_like_ptr);
+    }
+
+    if results.is_empty() {
+        return None;
+    }
+    Some(results)
+}
+
+fn unqualified_star_table_ptrs(
+    db: &dyn Db,
+    file: File,
+    target: &ast::Target,
+) -> Option<Vec<SyntaxNodePtr>> {
+    target.star_token()?;
+
+    let path = match ast_nav::target_parent_query(target.clone())? {
+        ast_nav::ParentQuery::Select(select) => {
+            let from_clause = select.from_clause()?;
+            let results = resolve::table_ptrs_from_clause(db, file, &from_clause);
+            if results.is_empty() {
+                return None;
+            }
+            return Some(results);
+        }
+        ast_nav::ParentQuery::Update(update) => update.relation_name()?.path(),
+        ast_nav::ParentQuery::Insert(insert) => insert.path(),
+        ast_nav::ParentQuery::Delete(delete) => delete.relation_name()?.path(),
+        ast_nav::ParentQuery::Merge(merge) => merge.relation_name()?.path(),
+    }?;
+
+    let position = target.syntax().text_range().start();
+    table_or_view_or_cte_ptrs(db, file, position, &path)
+}
+
+fn unqualified_star_in_arg_list_ptrs(
+    db: &dyn Db,
+    file: File,
+    arg_list: &ast::ArgList,
+) -> Option<Vec<SyntaxNodePtr>> {
+    let from_clause = arg_list
+        .syntax()
+        .ancestors()
+        .find_map(ast::Select::cast)?
+        .from_clause()?;
+    let results = resolve::table_ptrs_from_clause(db, file, &from_clause);
+
+    if results.is_empty() {
+        return None;
+    }
+
+    Some(results)
+}
+
 #[cfg(test)]
 mod test {
     use crate::db::{Database, File};