Fix issue when using CTE and UNION, refs tidyverse#1557

aguynamedryan · Nov 21, 2024 · efa8fe6 · efa8fe6
1 parent d854986
commit efa8fe6
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 1 deletion.
diff --git a/R/query-set-op.R b/R/query-set-op.R
@@ -110,7 +110,7 @@ flatten_query.union_query <- function(qry, query_list, con) {
   query_list_new <- flatten_query(x, query_list, con)
   qry$x <- get_subquery_name(x, query_list_new)
 
-  for (i in vctrs::vec_seq_along(qry$unions)) {
+  for (i in vctrs::vec_seq_along(qry$unions$table)) {
     y <- qry$unions$table[[i]]
     query_list_new <- flatten_query(y, query_list_new, con)
     qry$unions$table[[i]] <- get_subquery_name(y, query_list_new)

diff --git a/tests/testthat/_snaps/verb-set-ops.md b/tests/testthat/_snaps/verb-set-ops.md
@@ -62,3 +62,21 @@
       LEFT JOIN `lf1`
         ON (`LHS`.`x` = `lf1`.`x`)
 
+---
+
+    Code
+      lf1 %>% union_all(lf2) %>% show_query(sql_options = sql_options(cte = TRUE))
+    Output
+      <SQL>
+      WITH `q01` AS (
+        SELECT NULL AS `x`, `lf2`.*
+        FROM `lf2`
+      )
+      SELECT *
+      FROM `lf1`
+      
+      UNION ALL
+      
+      SELECT *
+      FROM `q01`
+
diff --git a/tests/testthat/test-verb-set-ops.R b/tests/testthat/test-verb-set-ops.R
@@ -75,6 +75,13 @@ test_that("can combine multiple union in one query", {
       show_query(sql_options = sql_options(cte = TRUE))
   )
 
+  # cte works with simple union
+  expect_snapshot(
+    lf1 %>%
+      union_all(lf2) %>%
+      show_query(sql_options = sql_options(cte = TRUE))
+  )
+
   lf_union <- lf1 %>%
     union_all(lf2) %>%
     union(lf3)