ide: support inherits/like tables in inlay hints

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("

crates/squawk_ide/src/resolve.rs

@@ -1737,6 +1737,80 @@ pub(crate) fn extract_column_name(col: &ast::Column) -> Option<Name> {
     Some(name)
 }
 
+pub(crate) fn collect_columns_from_create_table(
+    binder: &Binder,
+    root: &SyntaxNode,
+    create_table: &ast::CreateTableLike,
+) -> Vec<(Name, Option<SyntaxNodePtr>)> {
+    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<SyntaxNodePtr>)>,
+    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,