ide: goto def with insert/delete & CTE (#794)

Author: sbdchd

crates/squawk_ide/src/goto_definition.rs

@@ -3002,4 +3002,79 @@ select messages.message$0 from users left join messages on users.id = messages.u
           ╰╴                      ─ 1. source
         ");
     }
+
+    #[test]
+    fn goto_insert_select_cte_column() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+insert into users (id, email)
+select id$0, email from new_data;
+"), @r"
+          ╭▸ 
+        4 │     select 1 as id, '[email protected]' as email
+          │                 ── 2. destination
+          ‡
+        7 │ select id, email from new_data;
+          ╰╴        ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_insert_select_cte_column_second() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+insert into users (id, email)
+select id, email$0 from new_data;
+"), @r"
+          ╭▸ 
+        4 │     select 1 as id, '[email protected]' as email
+          │                                           ───── 2. destination
+          ‡
+        7 │ select id, email from new_data;
+          ╰╴               ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_insert_select_cte_table() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with new_data as (
+    select 1 as id, '[email protected]' as email
+)
+insert into users (id, email)
+select id, email from new_data$0;
+"), @r"
+          ╭▸ 
+        3 │ with new_data as (
+          │      ──────── 2. destination
+          ‡
+        7 │ select id, email from new_data;
+          ╰╴                             ─ 1. source
+        ");
+    }
+
+    #[test]
+    fn goto_delete_cte_column() {
+        assert_snapshot!(goto("
+create table users(id int, email text);
+with old_data as (
+    select 1 as id
+)
+delete from users where id in (select id$0 from old_data);
+"), @r"
+          ╭▸ 
+        4 │     select 1 as id
+          │                 ── 2. destination
+        5 │ )
+        6 │ delete from users where id in (select id from old_data);
+          ╰╴                                       ─ 1. source
+        ");
+    }
 }

crates/squawk_ide/src/resolve.rs

@@ -1,6 +1,6 @@
 use rowan::{TextRange, TextSize};
 use squawk_syntax::{
-    SyntaxNodePtr,
+    SyntaxNode, SyntaxNodePtr,
     ast::{self, AstNode},
 };
 
@@ -1120,9 +1120,7 @@ pub(crate) fn find_column_in_table(
 }
 
 fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option<SyntaxNodePtr> {
-    let select = name_ref.syntax().ancestors().find_map(ast::Select::cast)?;
-    let with_clause = select.with_clause()?;
-
+    let with_clause = find_parent_with_clause(name_ref.syntax())?;
     for with_table in with_clause.with_tables() {
         if let Some(name) = with_table.name()
             && Name::from_node(&name) == *cte_name
@@ -1134,17 +1132,26 @@ fn resolve_cte_table(name_ref: &ast::NameRef, cte_name: &Name) -> Option<SyntaxN
     None
 }
 
+fn find_parent_with_clause(node: &SyntaxNode) -> Option<ast::WithClause> {
+    node.ancestors().find_map(|x| {
+        if let Some(select) = ast::Select::cast(x.clone()) {
+            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) {
+            insert.with_clause()
+        } else {
+            None
+        }
+    })
+}
+
 fn resolve_cte_column(
     name_ref: &ast::NameRef,
     cte_name: &Name,
     column_name: &Name,
 ) -> Option<SyntaxNodePtr> {
-    let select = name_ref
-        .syntax()
-        .ancestors()
-        .filter_map(ast::Select::cast)
-        .find(|s| s.with_clause().is_some())?;
-    let with_clause = select.with_clause()?;
+    let with_clause = find_parent_with_clause(name_ref.syntax())?;
 
     for with_table in with_clause.with_tables() {
         if let Some(name) = with_table.name()

crates/squawk_syntax/src/ast/generated/nodes.rs

@@ -1263,6 +1263,7 @@ Table =
   'table' RelationName
 
 Insert =
+  WithClause?
   'insert' 'into' Path Alias? ColumnList?
   ('overriding' ('system' | 'user') 'value')?
   ('default' 'values' | Values | Stmt)
@@ -1326,6 +1327,7 @@ SetExpr =
 | 'default'
 
 Update =
+  WithClause?
   'update'
   RelationName
   Alias?
@@ -1348,6 +1350,7 @@ ReturningOption =
   Name
 
 Delete =
+  WithClause?
   'delete' 'from' RelationName Alias?
   UsingClause?
   (WhereClause | WhereCurrentOf)?