Restore crud flags if visiting_cte is set

Fixed bug where a CTE of an INSERT/UPDATE/DELETE that also uses RETURNING
could then not be SELECTed from directly, as the internal state of the
compiler would try to treat the outer SELECT as a DELETE statement itself
and access nonexistent state.

Fixes: #5181
Change-Id: Icba76f2148c8344baa1c04bac4ab6c6d24f23072

1 files changed, 35 insertions, 0 deletions

diff --git a/test/sql/test_cte.py b/test/sql/test_cte.py
index 4a7a80e77..c9178d580 100644
--- a/test/sql/test_cte.py
+++ b/test/sql/test_cte.py
@@ -999,6 +999,8 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "upsert.quantity FROM upsert))",
         )
 
+        eq_(insert.compile().isinsert, True)
+
     def test_anon_update_cte(self):
         orders = table("orders", column("region"))
         stmt = (
@@ -1016,6 +1018,8 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "SELECT anon_1.region FROM anon_1",
         )
 
+        eq_(stmt.select().compile().isupdate, False)
+
     def test_anon_insert_cte(self):
         orders = table("orders", column("region"))
         stmt = (
@@ -1028,6 +1032,7 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "VALUES (:region) RETURNING orders.region) "
             "SELECT anon_1.region FROM anon_1",
         )
+        eq_(stmt.select().compile().isinsert, False)
 
     def test_pg_example_one(self):
         products = table("products", column("id"), column("date"))
@@ -1054,6 +1059,33 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "INSERT INTO products_log (id, date) "
             "SELECT moved_rows.id, moved_rows.date FROM moved_rows",
         )
+        eq_(stmt.compile().isinsert, True)
+        eq_(stmt.compile().isdelete, False)
+
+    def test_pg_example_one_select_only(self):
+        products = table("products", column("id"), column("date"))
+
+        moved_rows = (
+            products.delete()
+            .where(
+                and_(products.c.date >= "dateone", products.c.date < "datetwo")
+            )
+            .returning(*products.c)
+            .cte("moved_rows")
+        )
+
+        stmt = moved_rows.select()
+
+        self.assert_compile(
+            stmt,
+            "WITH moved_rows AS "
+            "(DELETE FROM products WHERE products.date >= :date_1 "
+            "AND products.date < :date_2 "
+            "RETURNING products.id, products.date) "
+            "SELECT moved_rows.id, moved_rows.date FROM moved_rows",
+        )
+
+        eq_(stmt.compile().isdelete, False)
 
     def test_pg_example_two(self):
         products = table("products", column("id"), column("price"))
@@ -1076,6 +1108,7 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "SELECT t.id, t.price "
             "FROM t",
         )
+        eq_(stmt.compile().isupdate, False)
 
     def test_pg_example_three(self):
 
@@ -1136,6 +1169,7 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "SELECT pd.id, pd.price "
             "FROM pd",
         )
+        eq_(stmt.compile().isinsert, False)
 
     def test_update_pulls_from_cte(self):
         products = table("products", column("id"), column("price"))
@@ -1154,6 +1188,7 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             "UPDATE products SET id=:id, price=:price FROM pd "
             "WHERE products.price = pd.price",
         )
+        eq_(stmt.compile().isupdate, True)
 
     def test_standalone_function(self):
         a = table("a", column("x"))