diff --git a/crates/squawk_ide/src/inlay_hints.rs b/crates/squawk_ide/src/inlay_hints.rs index e482840b..2ba34114 100644 --- a/crates/squawk_ide/src/inlay_hints.rs +++ b/crates/squawk_ide/src/inlay_hints.rs @@ -138,21 +138,12 @@ fn inlay_hint_insert( }) .collect() } else { - // TODO: this doesn't work when there's `inherit`/`like` involved in the table def // `insert into t values (1, 2, 3)` - create_table? - .table_arg_list()? - .args() - .filter_map(|arg| { - if let ast::TableArg::Column(column) = arg - && let Some(name) = column.name() - { - let col_name = Name::from_node(&name); - let target = Some(name.syntax().text_range()); - Some((col_name, target, location.as_ref().map(|x| x.file))) - } else { - None - } + resolve::collect_columns_from_create_table(&binder, file.syntax(), &create_table?) + .into_iter() + .map(|(col_name, ptr)| { + let target = ptr.map(|p| p.to_node(file.syntax()).text_range()); + (col_name, target, location.as_ref().map(|x| x.file)) }) .collect() }; @@ -448,6 +439,34 @@ insert into t values (1, 2, 3); "); } + #[test] + fn insert_table_inherits_select() { + assert_snapshot!(check_inlay_hints(" +create table t (a int, b int); +create table u (c int) inherits (t); +insert into u select 1, 2, 3; +"), @r" + inlay hints: + ╭▸ + 4 │ insert into u select a: 1, b: 2, c: 3; + ╰╴ ─── ─── ─── + "); + } + + #[test] + fn insert_table_like_select() { + assert_snapshot!(check_inlay_hints(" +create table x (a int, b int); +create table y (c int, like x); +insert into y select 1, 2, 3; +"), @r" + inlay hints: + ╭▸ + 4 │ insert into y select c: 1, a: 2, b: 3; + ╰╴ ─── ─── ─── + "); + } + #[test] fn insert_select() { assert_snapshot!(check_inlay_hints(" diff --git a/crates/squawk_ide/src/resolve.rs b/crates/squawk_ide/src/resolve.rs index c8b82386..a31b6eb0 100644 --- a/crates/squawk_ide/src/resolve.rs +++ b/crates/squawk_ide/src/resolve.rs @@ -1737,6 +1737,80 @@ pub(crate) fn extract_column_name(col: &ast::Column) -> Option { Some(name) } +pub(crate) fn collect_columns_from_create_table( + binder: &Binder, + root: &SyntaxNode, + create_table: &ast::CreateTableLike, +) -> Vec<(Name, Option)> { + let mut columns = vec![]; + collect_columns_from_create_table_impl(binder, root, create_table, &mut columns, 0); + columns +} + +fn collect_columns_from_create_table_impl( + binder: &Binder, + root: &SyntaxNode, + create_table: &ast::CreateTableLike, + columns: &mut Vec<(Name, Option)>, + depth: usize, +) { + if depth > 40 { + log::info!("max depth reached, probably in a cycle"); + return; + } + + if let Some(inherits) = create_table.inherits() { + for path in inherits.paths() { + if let Some((table_name, schema)) = extract_table_schema_from_path(&path) { + let position = path.syntax().text_range().start(); + if let Some(ResolvedTableName::Table(parent_table)) = + resolve_table_name(binder, root, &table_name, &schema, position) + { + collect_columns_from_create_table_impl( + binder, + root, + &parent_table, + columns, + depth + 1, + ); + } + } + } + } + + if let Some(arg_list) = create_table.table_arg_list() { + for arg in arg_list.args() { + match &arg { + ast::TableArg::Column(column) => { + if let Some(name) = column.name() { + let col_name = Name::from_node(&name); + columns.push((col_name, Some(SyntaxNodePtr::new(name.syntax())))); + } + } + ast::TableArg::LikeClause(like_clause) => { + if let Some(path) = like_clause.path() + && let Some((table_name, schema)) = extract_table_schema_from_path(&path) + { + let position = path.syntax().text_range().start(); + if let Some(ResolvedTableName::Table(source_table)) = + resolve_table_name(binder, root, &table_name, &schema, position) + { + collect_columns_from_create_table_impl( + binder, + root, + &source_table, + columns, + depth + 1, + ); + } + } + } + ast::TableArg::TableConstraint(_) => (), + } + } + } +} + pub(crate) fn find_column_in_create_table( binder: &Binder, root: &SyntaxNode,