From 7d8a27fb4ac86f9e66c021e58aadd37d82a455bf Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Sat, 30 May 2026 23:30:05 -0400 Subject: [PATCH] ide: treat select into like create table as --- crates/squawk_ide/src/ast_nav.rs | 5 + crates/squawk_ide/src/classify.rs | 3 + crates/squawk_ide/src/collect.rs | 71 ++++++++----- crates/squawk_ide/src/completion.rs | 38 +++++++ crates/squawk_ide/src/goto_definition.rs | 16 +++ crates/squawk_ide/src/hover.rs | 122 +++++++++++++++++++++++ crates/squawk_ide/src/resolve.rs | 18 +++- 7 files changed, 244 insertions(+), 29 deletions(-) diff --git a/crates/squawk_ide/src/ast_nav.rs b/crates/squawk_ide/src/ast_nav.rs index 7b47bebc..6ea6a3fc 100644 --- a/crates/squawk_ide/src/ast_nav.rs +++ b/crates/squawk_ide/src/ast_nav.rs @@ -207,6 +207,7 @@ pub(crate) enum ParentSouce { CreateTableAs(ast::CreateTableAs), CreateView(ast::CreateViewLike), ParenSelect(ast::ParenSelect), + SelectInto(ast::SelectInto), WithTable(ast::WithTable), } @@ -239,6 +240,10 @@ pub(crate) fn parent_source(node: &SyntaxNode) -> Option { if let Some(create_table) = ast::CreateTableLike::cast(ancestor.clone()) { return Some(ParentSouce::CreateTable(create_table)); } + + if let Some(select_into) = ast::SelectInto::cast(ancestor.clone()) { + return Some(ParentSouce::SelectInto(select_into)); + } } None diff --git a/crates/squawk_ide/src/classify.rs b/crates/squawk_ide/src/classify.rs index 91579d3e..6b3cdcae 100644 --- a/crates/squawk_ide/src/classify.rs +++ b/crates/squawk_ide/src/classify.rs @@ -848,6 +848,9 @@ pub(crate) fn classify_def_node(def_node: &SyntaxNode) -> Option { if ast::CreateTableAs::can_cast(ancestor.kind()) { return Some(LocationKind::Table); } + if ast::SelectInto::can_cast(ancestor.kind()) { + return Some(LocationKind::Table); + } if ast::CreateIndex::can_cast(ancestor.kind()) { return Some(LocationKind::Index); } diff --git a/crates/squawk_ide/src/collect.rs b/crates/squawk_ide/src/collect.rs index 903ff6c6..08acee6b 100644 --- a/crates/squawk_ide/src/collect.rs +++ b/crates/squawk_ide/src/collect.rs @@ -102,14 +102,12 @@ fn resolved_to_column_ptrs( .map(|(name, _ty)| (name, None)) .collect() } - ResolvedTableName::SelectInto(select_into) => select_into - .select_clause() - .and_then(|c| c.target_list()) - .map(|t| target_list_columns_with_types(&t)) - .unwrap_or_default() - .into_iter() - .map(|(name, _ty)| (name, None)) - .collect(), + ResolvedTableName::SelectInto(select_into) => { + select_into_columns_with_types(db, file, &select_into) + .into_iter() + .map(|(name, _ty)| (name, None)) + .collect() + } ResolvedTableName::View(create_view_like) => { view_like_columns_with_types(db, file, &create_view_like) .into_iter() @@ -199,11 +197,9 @@ fn resolved_to_columns_with_types( ResolvedTableName::TableAs(create_table_as) => { select_columns_with_types(db, file, &create_table_as.query()) } - ResolvedTableName::SelectInto(select_into) => select_into - .select_clause() - .and_then(|c| c.target_list()) - .map(|t| target_list_columns_with_types(&t)) - .unwrap_or_default(), + ResolvedTableName::SelectInto(select_into) => { + select_into_columns_with_types(db, file, &select_into) + } ResolvedTableName::View(create_view_like) => { view_like_columns_with_types(db, file, &create_view_like) } @@ -225,6 +221,27 @@ pub(crate) fn create_table_as_columns_with_types( vec![] } +pub(crate) fn select_into_columns_with_types( + db: &dyn Db, + file: File, + select_into: &ast::SelectInto, +) -> Vec<(Name, Option)> { + let Some(target_list) = select_into.select_clause().and_then(|c| c.target_list()) else { + return vec![]; + }; + let from_clause = select_into.from_clause(); + + for file in list_files(db, file) { + let columns = + target_list_columns_with_types_in_file(db, file, &target_list, from_clause.as_ref()); + if !columns.is_empty() { + return columns; + } + } + + vec![] +} + fn target_list_columns_with_types_in_file( db: &dyn Db, file: File, @@ -578,6 +595,12 @@ fn column_type_at_location(db: &dyn Db, def: Location) -> Option { .find(|(n, _)| *n == column_name) .and_then(|(_, t)| t) } + ast_nav::ParentSouce::SelectInto(select_into) => { + select_into_columns_with_types(db, def.file, &select_into) + .into_iter() + .find(|(n, _)| *n == column_name) + .and_then(|(_, t)| t) + } ast_nav::ParentSouce::Alias(alias) => { let from_item = alias.syntax().ancestors().find_map(ast::FromItem::cast)?; columns_for_star_from_alias(db, def.file, &from_item, &alias) @@ -712,21 +735,11 @@ fn columns_for_star_from_table_ptr( Some(ast_nav::ParentSouce::ParenSelect(paren_select)) => { paren_select_columns_with_types(db, file, &paren_select) } - None => vec![], - } -} - -fn target_list_columns_with_types(target_list: &ast::TargetList) -> Vec<(Name, Option)> { - let mut columns = vec![]; - for target in target_list.targets() { - if let Some((col_name, _node)) = ColumnName::from_target(target.clone()) - && let Some(col_name_str) = col_name.to_string() - { - let ty = target.expr().and_then(|e| infer_type_from_expr(&e)); - columns.push((Name::from_string(col_name_str), ty)); + Some(ast_nav::ParentSouce::SelectInto(select_into)) => { + select_into_columns_with_types(db, file, &select_into) } + None => vec![], } - columns } pub(crate) fn paren_select_columns_with_types( @@ -832,6 +845,12 @@ pub(crate) fn star_column_names(db: &dyn Db, file: File, table_ptr: &SyntaxNodeP } return star_column_names_from_paren_select(db, file, &paren_select); } + Some(ast_nav::ParentSouce::SelectInto(select_into)) => { + select_into_columns_with_types(db, file, &select_into) + .into_iter() + .map(|(name, _ty)| name) + .collect() + } None => vec![], } } diff --git a/crates/squawk_ide/src/completion.rs b/crates/squawk_ide/src/completion.rs index e6798842..4ad6d592 100644 --- a/crates/squawk_ide/src/completion.rs +++ b/crates/squawk_ide/src/completion.rs @@ -448,6 +448,18 @@ fn column_completions_from_clause( sort_text: None, })); } + Some(ast_nav::ParentSouce::SelectInto(select_into)) => { + let columns = collect::select_into_columns_with_types(db, file, &select_into); + completions.extend(columns.into_iter().map(|(name, ty)| CompletionItem { + label: name.to_string(), + kind: CompletionItemKind::Column, + detail: ty.map(|t| t.to_string()), + insert_text: None, + insert_text_format: None, + trigger_completion_after_insert: false, + sort_text: None, + })); + } None => {} } } @@ -501,6 +513,12 @@ fn alias_base_columns_with_types( .map(|(name, ty)| (name, ty.map(|t| t.to_string()))) .collect() } + Some(ast_nav::ParentSouce::SelectInto(select_into)) => { + collect::select_into_columns_with_types(db, file, &select_into) + .into_iter() + .map(|(name, ty)| (name, ty.map(|t| t.to_string()))) + .collect() + } Some(ast_nav::ParentSouce::Alias(_)) | None => vec![], } } @@ -1207,6 +1225,26 @@ select $0 from child; "); } + #[test] + fn completion_after_select_select_into() { + assert_snapshot!(completions(" +select 1 a, 'x'::text b into t; +select $0 from t; +"), @" + label | kind | detail + --------------------+----------+--------- + a | Column | integer + b | Column | text + t | Table | + * | Operator | + public | Schema | + pg_catalog | Schema | + pg_temp | Schema | + pg_toast | Schema | + information_schema | Schema | + "); + } + #[test] fn completion_select_table_qualified() { assert_snapshot!(completions(" diff --git a/crates/squawk_ide/src/goto_definition.rs b/crates/squawk_ide/src/goto_definition.rs index 40d3ba9f..b7f60dc1 100644 --- a/crates/squawk_ide/src/goto_definition.rs +++ b/crates/squawk_ide/src/goto_definition.rs @@ -3049,6 +3049,22 @@ select a from t$0; "); } + #[test] + fn goto_select_into_select_star() { + assert_snapshot!(goto(" +create table t(a bigint); +select * into u from t; +select a$0 from u; +"), @" + ╭▸ + 2 │ create table t(a bigint); + │ ─ 2. destination + 3 │ select * into u from t; + 4 │ select a from u; + ╰╴ ─ 1. source + "); + } + #[test] fn goto_create_table_as_column() { assert_snapshot!(goto(" diff --git a/crates/squawk_ide/src/hover.rs b/crates/squawk_ide/src/hover.rs index 9ef33294..a3286d66 100644 --- a/crates/squawk_ide/src/hover.rs +++ b/crates/squawk_ide/src/hover.rs @@ -520,6 +520,32 @@ fn format_hover_for_column_ptr(db: &dyn Db, def: Location) -> Option { def_node, )); } + ast_nav::ParentSouce::SelectInto(select_into) => { + let column_name = collect::column_name_from_node(def_node)?; + let path = select_into.into_clause()?.path()?; + let (schema, table_name) = + resolve::resolve_table_info(db, InFile::new(def.file, &path))?; + let ty = collect::select_into_columns_with_types(db, def.file, &select_into) + .into_iter() + .find(|(name, _)| *name == column_name) + .and_then(|(_, ty)| ty); + return Some(hover_column_with_preceding_comment( + match ty { + Some(ty) => ColumnHover::schema_table_column_type( + &schema.to_string(), + &table_name, + &column_name.to_string(), + &ty.to_string(), + ), + None => ColumnHover::schema_table_column( + &schema.to_string(), + &table_name, + &column_name.to_string(), + ), + }, + def_node, + )); + } ast_nav::ParentSouce::CreateTable(create_table) => { let column = def_node.ancestors().find_map(ast::Column::cast)?; let column_name = column.name()?; @@ -600,6 +626,9 @@ fn format_table_source(db: &dyn Db, source: InFile) -> Opt format_create_table_as(db, InFile::new(file, create_table_as)) } ast_nav::ParentSouce::ParenSelect(paren_select) => format_paren_select(paren_select), + ast_nav::ParentSouce::SelectInto(select_into) => { + format_select_into(db, InFile::new(file, select_into)) + } } } @@ -749,6 +778,9 @@ fn hover_qualified_star_columns( ast_nav::ParentSouce::ParenSelect(paren_select) => { hover_qualified_star_columns_from_subquery(db, InFile::new(file, &paren_select)) } + ast_nav::ParentSouce::SelectInto(select_into) => { + hover_qualified_star_columns_from_select_into(db, InFile::new(file, &select_into)) + } } } @@ -844,6 +876,39 @@ fn hover_qualified_star_columns_from_table_as( merge_hovers(results) } +fn hover_qualified_star_columns_from_select_into( + db: &dyn Db, + select_into: InFile<&ast::SelectInto>, +) -> Option { + let file = select_into.file_id; + let select_into = select_into.value; + let path = select_into.into_clause()?.path()?; + let (schema, table_name) = resolve::resolve_table_info(db, InFile::new(file, &path))?; + let schema_str = schema.to_string(); + + let columns = collect::select_into_columns_with_types(db, file, select_into); + let results: Vec = columns + .into_iter() + .map(|(column_name, ty)| { + if let Some(ty) = ty { + return Hover::snippet(ColumnHover::schema_table_column_type( + &schema_str, + &table_name, + &column_name.to_string(), + &ty.to_string(), + )); + } + Hover::snippet(ColumnHover::schema_table_column( + &schema_str, + &table_name, + &column_name.to_string(), + )) + }) + .collect(); + + merge_hovers(results) +} + fn hover_qualified_star_columns_from_cte( db: &dyn Db, with_table: InFile, @@ -1211,6 +1276,14 @@ fn format_create_table_as( ))) } +fn format_select_into(db: &dyn Db, select_into: InFile) -> Option { + let file = select_into.file_id; + let select_into = select_into.value; + let path = select_into.into_clause()?.path()?; + let (schema, table_name) = resolve::resolve_table_info(db, InFile::new(file, &path))?; + Some(Hover::snippet(format!("table {schema}.{table_name}"))) +} + fn format_create_view(db: &dyn Db, def: Location) -> Option { let create_view = ast::CreateViewLike::cast(def.to_node(db)?)?; format_create_view_like(db, InFile::new(def.file, create_view)) @@ -5018,6 +5091,55 @@ select *$0 from child; "); } + #[test] + fn hover_select_into_column() { + assert_snapshot!(check_hover(" +select 1 a into t; +select a$0 from t; +"), @" + hover: column public.t.a integer + ╭▸ + 3 │ select a from t; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_into_star() { + assert_snapshot!(check_hover_info(" +select 1 a, 'x'::text b into t; +select *$0 from t; +").snippet, @" + column public.t.a integer + column public.t.b text + "); + } + + #[test] + fn hover_select_into_table() { + assert_snapshot!(check_hover(" +select 1 a into t; +select a from t$0; +"), @" + hover: table public.t + ╭▸ + 3 │ select a from t; + ╰╴ ─ hover + "); + } + + #[test] + fn hover_select_into_table_definition() { + assert_snapshot!(check_hover(" +select 1 a into t$0; +"), @" + hover: table public.t + ╭▸ + 2 │ select 1 a into t; + ╰╴ ─ hover + "); + } + #[test] fn hover_create_table_like_view_star() { assert_snapshot!(check_hover_info(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index 3694c343..536a808c 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1278,7 +1278,7 @@ fn resolve_column_for_table( return resolve_function(db, column_name, schemas, None, file); } ResolvedTableName::SelectInto(select_into) => { - if let Some(ptr) = find_column_in_select_into(file, &select_into, column_name) { + if let Some(ptr) = find_column_in_select_into(db, file, &select_into, column_name) { return Some(ptr); } return resolve_function(db, column_name, schemas, None, file); @@ -1585,7 +1585,8 @@ fn resolve_column_from_table_or_view_or_cte_impl( } if let Some(select_into) = node.ancestors().find_map(ast::SelectInto::cast) { - if let Some(cols) = find_column_in_select_into(file, &select_into, column_name) { + if let Some(cols) = find_column_in_select_into(db, file, &select_into, column_name) + { return Some(cols); } @@ -1901,7 +1902,7 @@ fn find_column_in_resolved( find_column_in_create_table_as(db, InFile::new(file, &create_table_as), column_name) } ResolvedTableName::SelectInto(select_into) => { - find_column_in_select_into(file, &select_into, column_name) + find_column_in_select_into(db, file, &select_into, column_name) } ResolvedTableName::View(create_view) => { find_column_in_create_view_like(file, &create_view, column_name) @@ -2047,11 +2048,22 @@ fn find_column_in_from_clause( } fn find_column_in_select_into( + db: &dyn Db, file: File, select_into: &ast::SelectInto, column_name: &Name, ) -> Option> { + let from_clause = select_into.from_clause(); for target in select_into.select_clause()?.target_list()?.targets() { + if target.star_token().is_some() { + if let Some(from_clause) = &from_clause + && let Some(result) = + find_column_in_from_clause(db, InFile::new(file, from_clause), column_name) + { + return Some(result); + } + continue; + } if let Some((col_name, node)) = ColumnName::from_target(target.clone()) && let Some(col_name_str) = col_name.to_string() && Name::from_string(col_name_str) == *column_name