- step 1 - add cte(), some docs, a test, get to the point where

we see that the "self.isupdate" thing is conflicting

4 files changed, 205 insertions, 126 deletions

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index cc9a49a91..43276ac4f 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -1972,8 +1972,8 @@ class SQLCompiler(Compiled):
         MySQL overrides this.
 
         """
-        return from_table._compiler_dispatch(self, asfrom=True,
-                                             iscrud=True, **kw)
+        kw['asfrom'] = True
+        return from_table._compiler_dispatch(self, iscrud=True, **kw)
 
     def update_from_clause(self, update_stmt,
                            from_table, extra_froms,
diff --git a/lib/sqlalchemy/sql/dml.py b/lib/sqlalchemy/sql/dml.py
index 7b506f9db..a4a362c22 100644
--- a/lib/sqlalchemy/sql/dml.py
+++ b/lib/sqlalchemy/sql/dml.py
@@ -9,15 +9,18 @@ Provide :class:`.Insert`, :class:`.Update` and :class:`.Delete`.
 
 """
 
-from .base import Executable, _generative, _from_objects, DialectKWArgs
+from .base import Executable, _generative, _from_objects, DialectKWArgs, \
+    ColumnCollection
 from .elements import ClauseElement, _literal_as_text, Null, and_, _clone, \
     _column_as_key
-from .selectable import _interpret_as_from, _interpret_as_select, HasPrefixes
+from .selectable import _interpret_as_from, _interpret_as_select, \
+    HasPrefixes, HasCTE
 from .. import util
 from .. import exc
 
 
-class UpdateBase(DialectKWArgs, HasPrefixes, Executable, ClauseElement):
+class UpdateBase(
+        HasCTE, DialectKWArgs, HasPrefixes, Executable, ClauseElement):
     """Form the base for ``INSERT``, ``UPDATE``, and ``DELETE`` statements.
 
     """
@@ -141,6 +144,16 @@ class UpdateBase(DialectKWArgs, HasPrefixes, Executable, ClauseElement):
         """
         self._returning = cols
 
+    @util.memoized_property
+    def columns(self):
+        """A select-style column collection from the RETURNING
+        expression of this clause."""
+
+        self._columns = ColumnCollection()
+        for col in self._returning:
+            col._make_proxy(self)
+        return self._columns.as_immutable()
+
     @_generative
     def with_hint(self, text, selectable=None, dialect_name="*"):
         """Add a table hint for a single table to this
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index c3906c2f2..95ed10524 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -1223,6 +1223,158 @@ class CTE(Generative, HasSuffixes, Alias):
         )
 
 
+class HasCTE(object):
+    def cte(self, name=None, recursive=False):
+        """Return a new :class:`.CTE`, or Common Table Expression instance.
+
+        Common table expressions are a SQL standard whereby SELECT
+        statements can draw upon secondary statements specified along
+        with the primary statement, using a clause called "WITH".
+        Special semantics regarding UNION can also be employed to
+        allow "recursive" queries, where a SELECT statement can draw
+        upon the set of rows that have previously been selected.
+
+        CTEs can also be applied to DML constructs UPDATE, INSERT
+        and DELETE on some databases, both as a source of CTE rows
+        when combined with RETURNING, as well as a consumer of
+        CTE rows.
+
+        SQLAlchemy detects :class:`.CTE` objects, which are treated
+        similarly to :class:`.Alias` objects, as special elements
+        to be delivered to the FROM clause of the statement as well
+        as to a WITH clause at the top of the statement.
+
+        .. versionadded:: 0.7.6
+
+        .. versionchanged:: 1.1 Added support for UPDATE/INSERT/DELETE as
+           CTE, CTEs added to UPDATE/INSERT/DELETE.
+
+        :param name: name given to the common table expression.  Like
+         :meth:`._FromClause.alias`, the name can be left as ``None``
+         in which case an anonymous symbol will be used at query
+         compile time.
+        :param recursive: if ``True``, will render ``WITH RECURSIVE``.
+         A recursive common table expression is intended to be used in
+         conjunction with UNION ALL in order to derive rows
+         from those already selected.
+
+        The following examples illustrate two examples from
+        Postgresql's documentation at
+        http://www.postgresql.org/docs/8.4/static/queries-with.html.
+
+        Example 1, non recursive::
+
+            from sqlalchemy import (Table, Column, String, Integer,
+                                    MetaData, select, func)
+
+            metadata = MetaData()
+
+            orders = Table('orders', metadata,
+                Column('region', String),
+                Column('amount', Integer),
+                Column('product', String),
+                Column('quantity', Integer)
+            )
+
+            regional_sales = select([
+                                orders.c.region,
+                                func.sum(orders.c.amount).label('total_sales')
+                            ]).group_by(orders.c.region).cte("regional_sales")
+
+
+            top_regions = select([regional_sales.c.region]).\\
+                    where(
+                        regional_sales.c.total_sales >
+                        select([
+                            func.sum(regional_sales.c.total_sales)/10
+                        ])
+                    ).cte("top_regions")
+
+            statement = select([
+                        orders.c.region,
+                        orders.c.product,
+                        func.sum(orders.c.quantity).label("product_units"),
+                        func.sum(orders.c.amount).label("product_sales")
+                ]).where(orders.c.region.in_(
+                    select([top_regions.c.region])
+                )).group_by(orders.c.region, orders.c.product)
+
+            result = conn.execute(statement).fetchall()
+
+        Example 2, WITH RECURSIVE::
+
+            from sqlalchemy import (Table, Column, String, Integer,
+                                    MetaData, select, func)
+
+            metadata = MetaData()
+
+            parts = Table('parts', metadata,
+                Column('part', String),
+                Column('sub_part', String),
+                Column('quantity', Integer),
+            )
+
+            included_parts = select([
+                                parts.c.sub_part,
+                                parts.c.part,
+                                parts.c.quantity]).\\
+                                where(parts.c.part=='our part').\\
+                                cte(recursive=True)
+
+
+            incl_alias = included_parts.alias()
+            parts_alias = parts.alias()
+            included_parts = included_parts.union_all(
+                select([
+                    parts_alias.c.sub_part,
+                    parts_alias.c.part,
+                    parts_alias.c.quantity
+                ]).
+                    where(parts_alias.c.part==incl_alias.c.sub_part)
+            )
+
+            statement = select([
+                        included_parts.c.sub_part,
+                        func.sum(included_parts.c.quantity).
+                          label('total_quantity')
+                    ]).\\
+                    group_by(included_parts.c.sub_part)
+
+            result = conn.execute(statement).fetchall()
+
+        Example 3, an upsert using UPDATE and INSERT with CTEs::
+
+            orders = table(
+                'orders',
+                column('region'),
+                column('amount'),
+                column('product'),
+                column('quantity')
+            )
+
+            upsert = (
+                orders.update()
+                .where(orders.c.region == 'Region1')
+                .values(amount=1.0, product='Product1', quantity=1)
+                .returning(*(orders.c._all_columns)).cte('upsert'))
+
+            insert = orders.insert().from_select(
+                orders.c.keys(),
+                select([
+                    literal('Region1'), literal(1.0),
+                    literal('Product1'), literal(1)
+                ).where(exists(upsert.select()))
+            )
+
+        .. seealso::
+
+            :meth:`.orm.query.Query.cte` - ORM version of
+            :meth:`.SelectBase.cte`.
+
+        """
+        return CTE(self, name=name, recursive=recursive)
+
+
 class FromGrouping(FromClause):
     """Represent a grouping of a FROM clause"""
     __visit_name__ = 'grouping'
@@ -1497,7 +1649,7 @@ class ForUpdateArg(ClauseElement):
             self.of = None
 
 
-class SelectBase(Executable, FromClause):
+class SelectBase(HasCTE, Executable, FromClause):
     """Base class for SELECT statements.
 
 
@@ -1531,125 +1683,6 @@ class SelectBase(Executable, FromClause):
         """
         return self.as_scalar().label(name)
 
-    def cte(self, name=None, recursive=False):
-        """Return a new :class:`.CTE`, or Common Table Expression instance.
-
-        Common table expressions are a SQL standard whereby SELECT
-        statements can draw upon secondary statements specified along
-        with the primary statement, using a clause called "WITH".
-        Special semantics regarding UNION can also be employed to
-        allow "recursive" queries, where a SELECT statement can draw
-        upon the set of rows that have previously been selected.
-
-        SQLAlchemy detects :class:`.CTE` objects, which are treated
-        similarly to :class:`.Alias` objects, as special elements
-        to be delivered to the FROM clause of the statement as well
-        as to a WITH clause at the top of the statement.
-
-        .. versionadded:: 0.7.6
-
-        :param name: name given to the common table expression.  Like
-         :meth:`._FromClause.alias`, the name can be left as ``None``
-         in which case an anonymous symbol will be used at query
-         compile time.
-        :param recursive: if ``True``, will render ``WITH RECURSIVE``.
-         A recursive common table expression is intended to be used in
-         conjunction with UNION ALL in order to derive rows
-         from those already selected.
-
-        The following examples illustrate two examples from
-        Postgresql's documentation at
-        http://www.postgresql.org/docs/8.4/static/queries-with.html.
-
-        Example 1, non recursive::
-
-            from sqlalchemy import (Table, Column, String, Integer,
-                                    MetaData, select, func)
-
-            metadata = MetaData()
-
-            orders = Table('orders', metadata,
-                Column('region', String),
-                Column('amount', Integer),
-                Column('product', String),
-                Column('quantity', Integer)
-            )
-
-            regional_sales = select([
-                                orders.c.region,
-                                func.sum(orders.c.amount).label('total_sales')
-                            ]).group_by(orders.c.region).cte("regional_sales")
-
-
-            top_regions = select([regional_sales.c.region]).\\
-                    where(
-                        regional_sales.c.total_sales >
-                        select([
-                            func.sum(regional_sales.c.total_sales)/10
-                        ])
-                    ).cte("top_regions")
-
-            statement = select([
-                        orders.c.region,
-                        orders.c.product,
-                        func.sum(orders.c.quantity).label("product_units"),
-                        func.sum(orders.c.amount).label("product_sales")
-                ]).where(orders.c.region.in_(
-                    select([top_regions.c.region])
-                )).group_by(orders.c.region, orders.c.product)
-
-            result = conn.execute(statement).fetchall()
-
-        Example 2, WITH RECURSIVE::
-
-            from sqlalchemy import (Table, Column, String, Integer,
-                                    MetaData, select, func)
-
-            metadata = MetaData()
-
-            parts = Table('parts', metadata,
-                Column('part', String),
-                Column('sub_part', String),
-                Column('quantity', Integer),
-            )
-
-            included_parts = select([
-                                parts.c.sub_part,
-                                parts.c.part,
-                                parts.c.quantity]).\\
-                                where(parts.c.part=='our part').\\
-                                cte(recursive=True)
-
-
-            incl_alias = included_parts.alias()
-            parts_alias = parts.alias()
-            included_parts = included_parts.union_all(
-                select([
-                    parts_alias.c.sub_part,
-                    parts_alias.c.part,
-                    parts_alias.c.quantity
-                ]).
-                    where(parts_alias.c.part==incl_alias.c.sub_part)
-            )
-
-            statement = select([
-                        included_parts.c.sub_part,
-                        func.sum(included_parts.c.quantity).
-                          label('total_quantity')
-                    ]).\\
-                    group_by(included_parts.c.sub_part)
-
-            result = conn.execute(statement).fetchall()
-
-
-        .. seealso::
-
-            :meth:`.orm.query.Query.cte` - ORM version of
-            :meth:`.SelectBase.cte`.
-
-        """
-        return CTE(self, name=name, recursive=recursive)
-
     @_generative
     @util.deprecated('0.6',
                      message="``autocommit()`` is deprecated. Use "
diff --git a/test/sql/test_cte.py b/test/sql/test_cte.py
index b59914afc..885b15599 100644
--- a/test/sql/test_cte.py
+++ b/test/sql/test_cte.py
@@ -1,6 +1,6 @@
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing import AssertsCompiledSQL, assert_raises_message
-from sqlalchemy.sql import table, column, select, func, literal
+from sqlalchemy.sql import table, column, select, func, literal, exists
 from sqlalchemy.dialects import mssql
 from sqlalchemy.engine import default
 from sqlalchemy.exc import CompileError
@@ -492,3 +492,36 @@ class CTETest(fixtures.TestBase, AssertsCompiledSQL):
             'regional_sales WHERE "order"."order" > regional_sales."order"',
             dialect='postgresql'
         )
+
+    def test_upsert(self):
+        orders = table(
+            'orders',
+            column('region'),
+            column('amount'),
+            column('product'),
+            column('quantity')
+        )
+
+        upsert = (
+            orders.update()
+            .where(orders.c.region == 'Region1')
+            .values(amount=1.0, product='Product1', quantity=1)
+            .returning(*(orders.c._all_columns)).cte('upsert'))
+
+        insert = orders.insert().from_select(
+            orders.c.keys(),
+            select([
+                literal('Region1'), literal(1.0),
+                literal('Product1'), literal(1)
+            ]).where(exists(upsert.select()))
+        )
+
+        self.assert_compile(
+            insert,
+            "WITH upsert AS (UPDATE orders SET amount = 1.0, "
+            "product = 'Product1', quantity = 1 WHERE region = 'Region1' "
+            "RETURNING region, amount, product, quantity) "
+            "INSERT INTO orders (region, amount, product, quantity) "
+            "SELECT ('Region1', 1.0, 'Product1', 1) WHERE NOT EXISTS "
+            "(SELECT * FROM upsert)"
+        )