ide: hover/goto def for subquery w/ alias (#970)

Author: sbdchd

crates/squawk_ide/src/classify.rs

@@ -917,6 +917,9 @@ pub(crate) fn classify_def_node(def_node: &SyntaxNode) -> Option<NameRefClass> {
             return Some(NameRefClass::Channel);
         }
         if ast::Alias::can_cast(ancestor.kind()) {
+            if in_column {
+                return Some(NameRefClass::SelectColumn);
+            }
             return Some(NameRefClass::FromTable);
         }
         if ast::AsName::can_cast(ancestor.kind())

crates/squawk_ide/src/goto_definition.rs

@@ -4981,6 +4981,42 @@ select x, b$0 from (select 1 a, 2 b) t(x);
         ");
     }
 
+    #[test]
+    fn goto_subquery_alias_with_column_list_table_ref() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b)
+select z$0 from (select * from t) as z(x, y);
+"), @"
+          ╭▸ 
+        3 │ select z from (select * from t) as z(x, y);
+          ╰╴       ─ 1. source                 ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_subquery_alias_with_column_list_table_ref_shadows_column() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b)
+select z$0 from (select a as z, b from t) as z(x, y);
+"), @"
+          ╭▸ 
+        3 │ select z from (select a as z, b from t) as z(x, y);
+          ╰╴       ─ 1. source                         ─ 2. destination
+        ");
+    }
+
+    #[test]
+    fn goto_subquery_nested_paren_alias_with_column_list_table_ref() {
+        assert_snapshot!(goto("
+with t as (select 1 a, 2 b, 3 c)
+select z$0 from ((select * from t)) as z(x, y);
+"), @"
+          ╭▸ 
+        3 │ select z from ((select * from t)) as z(x, y);
+          ╰╴       ─ 1. source                   ─ 2. destination
+        ");
+    }
+
     #[test]
     fn goto_table_expr_values_cte_partial_alias() {
         assert_snapshot!(goto("

crates/squawk_ide/src/hover.rs

@@ -338,6 +338,18 @@ fn format_hover_for_column_node(
             return format_view_column(&create_view, column_name, binder);
         }
 
+        if let Some(alias) = ast::Alias::cast(a.clone())
+            && let Some(alias_name) = alias.name()
+            && alias.column_list().is_some()
+        {
+            let table_name = Name::from_node(&alias_name);
+            let column_name = Name::from_string(column_name_node.text().to_string());
+            return Some(ColumnHover::table_column(
+                &table_name.to_string(),
+                &column_name.to_string(),
+            ));
+        }
+
         if let Some(create_table_as) = ast::CreateTableAs::cast(a.clone()) {
             let column_name = if let Some(name) = ast::Name::cast(column_name_node.clone()) {
                 Name::from_node(&name)
@@ -581,6 +593,9 @@ fn hover_unqualified_star_in_arg_list(
 
 fn hover_subquery_table(def_node: &SyntaxNode) -> Option<String> {
     let alias = def_node.ancestors().find_map(ast::Alias::cast)?;
+    if alias.column_list().is_some() {
+        return None;
+    }
     let name = Name::from_node(&alias.name()?);
     let from_item = alias.syntax().ancestors().find_map(ast::FromItem::cast)?;
     let paren_select = from_item.paren_select()?;
@@ -671,10 +686,15 @@ fn collect_star_column_names(
     let table_name_node = table_ptr.to_node(root);
 
     if let Some(paren_select) = ast::ParenSelect::cast(table_name_node.clone()) {
-        return resolve::collect_paren_select_columns_with_types(binder, root, &paren_select)
-            .into_iter()
-            .map(|(name, _ty)| name)
-            .collect();
+        let columns: Vec<Name> =
+            resolve::collect_paren_select_columns_with_types(binder, root, &paren_select)
+                .into_iter()
+                .map(|(name, _ty)| name)
+                .collect();
+        if !columns.is_empty() {
+            return columns;
+        }
+        return collect_star_column_names_from_paren_select(root, &paren_select, binder);
     }
 
     match resolve::find_table_source(&table_name_node) {
@@ -720,6 +740,29 @@ fn collect_star_column_names(
     }
 }
 
+fn collect_star_column_names_from_paren_select(
+    root: &SyntaxNode,
+    paren_select: &ast::ParenSelect,
+    binder: &binder::Binder,
+) -> Vec<Name> {
+    let Some(select_variant) = paren_select.select() else {
+        return vec![];
+    };
+    let ast::SelectVariant::Select(select) = select_variant else {
+        return vec![];
+    };
+    let Some(from_clause) = select.from_clause() else {
+        return vec![];
+    };
+    let mut columns = vec![];
+    for from_item in from_clause.from_items() {
+        if let Some(table_ptr) = resolve::table_ptr_from_from_item(binder, &from_item) {
+            columns.extend(collect_star_column_names(root, &table_ptr, binder));
+        }
+    }
+    columns
+}
+
 fn hover_qualified_star_columns_from_table(
     root: &SyntaxNode,
     create_table: &impl ast::HasCreateTable,
@@ -3235,6 +3278,19 @@ select u$0.x, u.y from t as u(x, y);
         ");
     }
 
+    #[test]
+    fn hover_on_cte_table_alias_with_column_list_column_ref() {
+        assert_snapshot!(check_hover("
+with t as (select 1 a, 2 b, 3 c)
+select u.x$0 from t as u(x, y);
+"), @"
+        hover: column u.x
+          ╭▸ 
+        3 │ select u.x from t as u(x, y);
+          ╰╴         ─ hover
+        ");
+    }
+
     #[test]
     fn hover_on_cte_table_alias_with_column_list_table_ref() {
         assert_snapshot!(check_hover("
@@ -3248,6 +3304,32 @@ select u$0 from t as u(x, y);
         ");
     }
 
+    #[test]
+    fn hover_on_subquery_alias_with_column_list_table_ref() {
+        assert_snapshot!(check_hover("
+with t as (select 1 a, 2 b, 3 c)
+select z$0 from (select * from t) as z(x, y);
+"), @"
+        hover: table z(x, y, c)
+          ╭▸ 
+        3 │ select z from (select * from t) as z(x, y);
+          ╰╴       ─ hover
+        ");
+    }
+
+    #[test]
+    fn hover_on_subquery_nested_paren_alias_with_column_list_table_ref() {
+        assert_snapshot!(check_hover("
+with t as (select 1 a, 2 b, 3 c)
+select z$0 from ((select * from t)) as z(x, y);
+"), @"
+        hover: table z(x, y, c)
+          ╭▸ 
+        3 │ select z from ((select * from t)) as z(x, y);
+          ╰╴       ─ hover
+        ");
+    }
+
     #[test]
     fn hover_on_cte_table_alias_with_partial_column_list_star() {
         assert_snapshot!(check_hover("

crates/squawk_ide/src/resolve.rs

@@ -1285,14 +1285,23 @@ fn resolve_from_item_column_ptr(
     let column_name = Name::from_node(column_name_ref);
     if let Some(paren_select) = from_item.paren_select() {
         let alias = from_item.alias();
-        return resolve_subquery_column_ptr(
+        if let Some(ptr) = resolve_subquery_column_ptr(
             binder,
             root,
             &paren_select,
             column_name_ref,
             &column_name,
             alias.as_ref(),
-        );
+        ) {
+            return Some(ptr);
+        }
+        if let Some(alias) = alias
+            && let Some(alias_name) = alias.name()
+            && Name::from_node(&alias_name) == column_name
+        {
+            return Some(SyntaxNodePtr::new(alias_name.syntax()));
+        }
+        return None;
     }
 
     if let Some(paren_expr) = from_item.paren_expr() {
@@ -1308,13 +1317,18 @@ fn resolve_from_item_column_ptr(
             }
         }
 
-        return resolve_column_from_paren_expr(
-            binder,
-            root,
-            &paren_expr,
-            column_name_ref,
-            &column_name,
-        );
+        if let Some(ptr) =
+            resolve_column_from_paren_expr(binder, root, &paren_expr, column_name_ref, &column_name)
+        {
+            return Some(ptr);
+        }
+        if let Some(alias) = from_item.alias()
+            && let Some(alias_name) = alias.name()
+            && Name::from_node(&alias_name) == column_name
+        {
+            return Some(SyntaxNodePtr::new(alias_name.syntax()));
+        }
+        return None;
     }
 
     let alias_len = if let Some(alias) = from_item.alias()
@@ -2425,7 +2439,7 @@ fn resolve_subquery_column_ptr(
 ) -> Option<SyntaxNodePtr> {
     let select_variant = paren_select.select()?;
 
-    if let Some(alias) = alias
+    let column_list_len = if let Some(alias) = alias
         && let Some(column_list) = alias.column_list()
     {
         for col in column_list.columns() {
@@ -2438,7 +2452,10 @@ fn resolve_subquery_column_ptr(
         if matches!(select_variant, ast::SelectVariant::Values(_)) {
             return None;
         }
-    }
+        column_list.columns().count()
+    } else {
+        0
+    };
 
     // TODO: this should just be a match stmt
     if let ast::SelectVariant::Table(table) = select_variant {
@@ -2482,8 +2499,16 @@ fn resolve_subquery_column_ptr(
     let select_clause = subquery_select.select_clause()?;
     let target_list = select_clause.target_list()?;
 
-    for target in target_list.targets() {
+    for (i, target) in target_list.targets().enumerate() {
         if let Some((col_name, node)) = ColumnName::from_target(target.clone()) {
+            // skip targets that have been renamed by the column list
+            if i < column_list_len {
+                if matches!(col_name, ColumnName::Star) {
+                    // star expands to multiple columns, handled below
+                } else {
+                    continue;
+                }
+            }
             if let Some(col_name_str) = col_name.to_string()
                 && Name::from_string(col_name_str) == *column_name
             {
@@ -2759,6 +2784,10 @@ pub(crate) fn table_ptr_from_from_item(
         return Some(SyntaxNodePtr::new(paren_select.syntax()));
     }
 
+    if let Some(paren_expr) = from_item.paren_expr() {
+        return table_ptr_from_paren_expr(binder, &paren_expr);
+    }
+
     let (table_name, schema) = table_and_schema_from_from_item(from_item)?;
     let position = from_item.syntax().text_range().start();
 
@@ -2780,6 +2809,19 @@ pub(crate) fn table_ptr_from_from_item(
     None
 }
 
+fn table_ptr_from_paren_expr(
+    binder: &Binder,
+    paren_expr: &ast::ParenExpr,
+) -> Option<SyntaxNodePtr> {
+    if let Some(from_item) = paren_expr.from_item() {
+        return table_ptr_from_from_item(binder, &from_item);
+    }
+    if let Some(ast::Expr::ParenExpr(inner)) = paren_expr.expr() {
+        return table_ptr_from_paren_expr(binder, &inner);
+    }
+    None
+}
+
 fn collect_table_ptrs_from_join_expr(
     binder: &Binder,
     join_expr: &ast::JoinExpr,