ide: support ctes with update

crates/squawk_ide/src/goto_definition.rs

@@ -3252,4 +3252,76 @@ update users set email = messages.email from messages$0 where users.id = message
           ╰╴                                                    ─ 1. source
         ");
     }
+
+    #[test]
+    fn goto_update_with_cte_table() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+update users set email = new_data.email from new_data$0 where users.id = new_data.id;
+"), @r"
+          ╭▸ 
+        3 │ with new_data as (
+          │      ──────── 2. destination
+          ‡
+        6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
+          ╰╴                                                    ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_update_with_cte_column_in_set() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+update users set email = new_data.email$0 from new_data where users.id = new_data.id;
+"), @r"
+          ╭▸ 
+        4 │     select 1 as id, '[email protected]' as email
+          │                                          ───── 2. destination
+        5 │ )
+        6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
+          ╰╴                                      ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_update_with_cte_column_in_where() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+update users set email = new_data.email from new_data where new_data.id$0 = users.id;
+"), @r"
+          ╭▸ 
+        4 │     select 1 as id, '[email protected]' as email
+          │                 ── 2. destination
+        5 │ )
+        6 │ update users set email = new_data.email from new_data where new_data.id = users.id;
+          ╰╴                                                                      ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_update_with_cte_values() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    values (1, '[email protected]')
+)
+update users set email = new_data.column2$0 from new_data where users.id = new_data.column1;
+"), @r"
+          ╭▸ 
+        4 │     values (1, '[email protected]')
+          │                ───────────────── 2. destination
+        5 │ )
+        6 │ update users set email = new_data.column2 from new_data where users.id = new_data.column1;
+          ╰╴                                        ─ 1. source
+        ");
+    }
 }

crates/squawk_ide/src/hover.rs

@@ -2380,4 +2380,52 @@ update users set email = messages.email from public.messages$0 where users.id =
           ╰╴                                                           ─ hover
         ");
     }
+
+    #[test]
+    fn hover_on_update_with_cte_table() {
+        assert_snapshot!(check_hover("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+update users set email = new_data.email from new_data$0 where users.id = new_data.id;
+"), @r"
+        hover: with new_data as (select 1 as id, '[email protected]' as email)
+          ╭▸ 
+        6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
+          ╰╴                                                    ─ hover
+        ");
+    }
+
+    #[test]
+    fn hover_on_update_with_cte_column_in_set() {
+        assert_snapshot!(check_hover("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+update users set email = new_data.email$0 from new_data where users.id = new_data.id;
+"), @r"
+        hover: column new_data.email
+          ╭▸ 
+        6 │ update users set email = new_data.email from new_data where users.id = new_data.id;
+          ╰╴                                      ─ hover
+        ");
+    }
+
+    #[test]
+    fn hover_on_update_with_cte_column_in_where() {
+        assert_snapshot!(check_hover("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+update users set email = new_data.email from new_data where new_data.id$0 = users.id;
+"), @r"
+        hover: column new_data.id
+          ╭▸ 
+        6 │ update users set email = new_data.email from new_data where new_data.id = users.id;
+          ╰╴                                                                      ─ hover
+        ");
+    }
 }

crates/squawk_ide/src/inlay_hints.rs

@@ -85,9 +85,10 @@ fn inlay_hint_insert(
 ) -> Option<()> {
     let values = insert.values()?;
     let row_list = values.row_list()?;
+    let create_table = resolve::resolve_insert_create_table(file, binder, &insert);
 
     let columns: Vec<(Name, Option<TextRange>)> = if let Some(column_list) = insert.column_list() {
-        let create_table = resolve::resolve_insert_create_table(file, binder, &insert);
+        // `insert into t(a, b, c) values (1, 2, 3)`
         column_list
             .columns()
             .filter_map(|col| {
@@ -100,8 +101,8 @@ fn inlay_hint_insert(
             })
             .collect()
     } else {
-        let create_table = resolve::resolve_insert_create_table(file, binder, &insert)?;
-        create_table
+        // `insert into t values (1, 2, 3)`
+        create_table?
             .table_arg_list()?
             .args()
             .filter_map(|arg| {

crates/squawk_ide/src/resolve.rs

@@ -925,6 +925,27 @@ fn resolve_update_where_column(binder: &Binder, name_ref: &ast::NameRef) -> Opti
     let column_name = Name::from_node(name_ref);
 
     let update = name_ref.syntax().ancestors().find_map(ast::Update::cast)?;
+
+    // `update t set a = b from u`
+    if let Some(from_clause) = update.from_clause() {
+        for from_item in from_clause.from_items() {
+            if let Some(result) = resolve_from_item_for_column(binder, &from_item, name_ref) {
+                return Some(result);
+            }
+        }
+
+        for join_expr in from_clause.join_exprs() {
+            if let Some(result) =
+                resolve_from_join_expr(&join_expr, &|from_item: &ast::FromItem| {
+                    resolve_from_item_for_column(binder, from_item, name_ref)
+                })
+            {
+                return Some(result);
+            }
+        }
+    }
+
+    // `update t set a = b`
     let relation_name = update.relation_name()?;
     let path = relation_name.path()?;
 
@@ -1131,8 +1152,10 @@ fn find_parent_with_clause(node: &SyntaxNode) -> Option<ast::WithClause> {
             select.with_clause()
         } else if let Some(delete) = ast::Delete::cast(x.clone()) {
             delete.with_clause()
-        } else if let Some(insert) = ast::Insert::cast(x) {
+        } else if let Some(insert) = ast::Insert::cast(x.clone()) {
             insert.with_clause()
+        } else if let Some(update) = ast::Update::cast(x) {
+            update.with_clause()
         } else {
             None
         }