Work w/ prefetch even for selects, if present

Fixed bug in new CTE feature for update/insert/delete stated
as a CTE inside of an enclosing statement (typically SELECT) whereby
oninsert and onupdate values weren't called upon for the embedded
statement.

This is accomplished by consulting prefetch
for all statements.  The collection is also broken into
separate insert/update collections so that we don't need to
consult toplevel self.isinsert to determine if the prefetch
is for an insert or an update.  What we don't yet test for
are CTE combinations that have both insert/update in one
statement, though these should now work in theory provided
the underlying database supports such a statement.

Change-Id: I3b6a860e22c86743c91c56a7ec751ff706f66f64
Fixes: #3745

2 files changed, 95 insertions, 0 deletions

diff --git a/test/requirements.py b/test/requirements.py
index d31088e16..3c7a3fbb4 100644
--- a/test/requirements.py
+++ b/test/requirements.py
@@ -351,6 +351,14 @@ class DefaultRequirements(SuiteRequirements):
         return skip_if(exclude('mysql', '<', (4, 1, 1)), 'no subquery support')
 
     @property
+    def ctes(self):
+        """Target database supports CTEs"""
+
+        return only_if(
+            ['postgresql', 'mssql']
+        )
+
+    @property
     def mod_operator_as_percent_sign(self):
         """target database must use a plain percent '%' as the 'modulus'
         operator."""
diff --git a/test/sql/test_defaults.py b/test/sql/test_defaults.py
index db19e145b..57af1e536 100644
--- a/test/sql/test_defaults.py
+++ b/test/sql/test_defaults.py
@@ -539,6 +539,93 @@ class DefaultTest(fixtures.TestBase):
         eq_(55, l['col3'])
 
 
+class CTEDefaultTest(fixtures.TablesTest):
+    __requires__ = ('ctes',)
+    __backend__ = True
+
+    @classmethod
+    def define_tables(cls, metadata):
+        Table(
+            'q', metadata,
+            Column('x', Integer, default=2),
+            Column('y', Integer, onupdate=5),
+            Column('z', Integer)
+        )
+
+        Table(
+            'p', metadata,
+            Column('s', Integer),
+            Column('t', Integer),
+            Column('u', Integer, onupdate=1)
+        )
+
+    def _test_a_in_b(self, a, b):
+        q = self.tables.q
+        p = self.tables.p
+
+        with testing.db.connect() as conn:
+            if a == 'delete':
+                conn.execute(q.insert().values(y=10, z=1))
+                cte = q.delete().\
+                    where(q.c.z == 1).returning(q.c.z).cte('c')
+                expected = None
+            elif a == "insert":
+                cte = q.insert().values(z=1, y=10).returning(q.c.z).cte('c')
+                expected = (2, 10)
+            elif a == "update":
+                conn.execute(q.insert().values(x=5, y=10, z=1))
+                cte = q.update().\
+                    where(q.c.z == 1).values(x=7).returning(q.c.z).cte('c')
+                expected = (7, 5)
+            elif a == "select":
+                conn.execute(q.insert().values(x=5, y=10, z=1))
+                cte = sa.select([q.c.z]).cte('c')
+                expected = (5, 10)
+
+            if b == "select":
+                conn.execute(p.insert().values(s=1))
+                stmt = select([p.c.s, cte.c.z])
+            elif b == "insert":
+                sel = select([1, cte.c.z, ])
+                stmt = p.insert().from_select(['s', 't'], sel).returning(
+                    p.c.s, p.c.t)
+            elif b == "delete":
+                stmt = p.insert().values(s=1, t=cte.c.z).returning(
+                    p.c.s, cte.c.z)
+            elif b == "update":
+                conn.execute(p.insert().values(s=1))
+                stmt = p.update().values(t=5).\
+                    where(p.c.s == cte.c.z).\
+                    returning(p.c.u, cte.c.z)
+            eq_(
+                conn.execute(stmt).fetchall(),
+                [(1, 1)]
+            )
+
+            eq_(
+                conn.execute(select([q.c.x, q.c.y])).fetchone(),
+                expected
+            )
+
+    def test_update_in_select(self):
+        self._test_a_in_b("update", "select")
+
+    def test_delete_in_select(self):
+        self._test_a_in_b("update", "select")
+
+    def test_insert_in_select(self):
+        self._test_a_in_b("update", "select")
+
+    def test_select_in_update(self):
+        self._test_a_in_b("select", "update")
+
+    def test_select_in_insert(self):
+        self._test_a_in_b("select", "insert")
+
+    # TODO: updates / inserts can be run in one statement w/ CTE ?
+    # deletes?
+
+
 class PKDefaultTest(fixtures.TablesTest):
     __requires__ = ('subqueries',)
     __backend__ = True