From 7f87bb1714bdccd5b131e612af6ee9e4d93c7490 Mon Sep 17 00:00:00 2001
From: Steve Dignam <steve@dignam.xyz>
Date: Tue, 13 Jan 2026 23:36:39 -0500
Subject: [PATCH] ide: support select-to-values action in more places

---
 crates/squawk_ide/src/code_actions.rs | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/crates/squawk_ide/src/code_actions.rs b/crates/squawk_ide/src/code_actions.rs
index 748add6a..eae5e85f 100644
--- a/crates/squawk_ide/src/code_actions.rs
+++ b/crates/squawk_ide/src/code_actions.rs
@@ -662,12 +662,13 @@ fn rewrite_select_as_values(
 
 fn find_select_parent(token: SyntaxToken) -> Option<SelectContext> {
     let mut found_select = None;
+    let mut found_compound = None;
     for node in token.parent_ancestors() {
         if let Some(compound_select) = ast::CompoundSelect::cast(node.clone()) {
             if compound_select.union_token().is_some() && compound_select.all_token().is_some() {
-                return Some(SelectContext::Compound(compound_select));
+                found_compound = Some(SelectContext::Compound(compound_select));
             } else {
-                return None;
+                break;
             }
         }
         if found_select.is_none()
@@ -676,7 +677,7 @@ fn find_select_parent(token: SyntaxToken) -> Option<SelectContext> {
             found_select = Some(SelectContext::Single(select));
         }
     }
-    found_select
+    found_compound.or(found_select)
 }
 
 #[cfg(test)]
@@ -1631,6 +1632,17 @@ mod test {
         );
     }
 
+    #[test]
+    fn rewrite_select_as_values_multiple_rows_cursor_on_second_union() {
+        assert_snapshot!(
+            apply_code_action(
+                rewrite_select_as_values,
+                "select 1 as column1, 2 as column2 union all select 3, 4 union$0 all select 5, 6;"
+            ),
+            @"values (1, 2), (3, 4), (5, 6);"
+        );
+    }
+
     #[test]
     fn rewrite_select_as_values_single_column() {
         assert_snapshot!(