add support for returning results from inserts and updates for postgresql 8.2+. [ticket:797]

2 files changed, 61 insertions, 13 deletions

diff --git a/lib/sqlalchemy/databases/postgres.py b/lib/sqlalchemy/databases/postgres.py
index 305e8e831..03b3fd042 100644
--- a/lib/sqlalchemy/databases/postgres.py
+++ b/lib/sqlalchemy/databases/postgres.py
@@ -10,13 +10,20 @@ PostgreSQL supports partial indexes. To create them pass a posgres_where
 option to the Index constructor::
 
   Index('my_index', my_table.c.id, postgres_where=tbl.c.value > 10)
+
+PostgreSQL 8.2+ supports returning a result set from inserts and updates.
+To use this pass the column/expression list to the postgres_returning
+parameter when creating the queries::
+    
+  raises = tbl.update(empl.c.sales > 100, values=dict(salary=empl.c.salary * 1.1), 
+    postgres_returning=[empl.c.id, empl.c.salary]).execute().fetchall()
 """
 
 import re, random, warnings, string
 
 from sqlalchemy import sql, schema, exceptions, util
 from sqlalchemy.engine import base, default
-from sqlalchemy.sql import compiler
+from sqlalchemy.sql import compiler, expression
 from sqlalchemy.sql import operators as sql_operators
 from sqlalchemy import types as sqltypes
 
@@ -198,13 +205,27 @@ def descriptor():
     ]}
 
 SELECT_RE = re.compile(
-    r'\s*(?:SELECT|FETCH)',
+    r'\s*(?:SELECT|FETCH|(UPDATE|INSERT))',
+    re.I | re.UNICODE)
+
+RETURNING_RE = re.compile(
+    'RETURNING',
+    re.I | re.UNICODE)
+
+# This finds if the RETURNING is not inside a quoted/commented values. Handles string literals,
+# quoted identifiers, dollar quotes, SQL comments and C style multiline comments. This does not
+# handle correctly nested C style quotes, lets hope no one does the following:
+# UPDATE tbl SET x=y /* foo /* bar */ RETURNING */
+RETURNING_QUOTED_RE = re.compile(
+    '\\s*(?:UPDATE|INSERT)\\s(?:[^\'"$/-]|-(?!-)|/(?!\\*)|"(?:[^"]|"")*"|\'(?:[^\']|\'\')*\'|\\$(?P<dquote>[^$]*)\\$.*?\\$(?P=dquote)\\$|--[^\n]*\n|/\\*([^*]|\\*(?!/))*\\*/)*\\sRETURNING',
     re.I | re.UNICODE)
 
 class PGExecutionContext(default.DefaultExecutionContext):
 
     def is_select(self):
-        return SELECT_RE.match(self.statement)
+        m = SELECT_RE.match(self.statement)
+        return m and (not m.group(1) or (RETURNING_RE.search(self.statement)
+           and RETURNING_QUOTED_RE.match(self.statement)))
         
     def create_cursor(self):
         # executing a default or Sequence standalone creates an execution context without a statement.  
@@ -598,6 +619,29 @@ class PGCompiler(compiler.DefaultCompiler):
         else:
             return super(PGCompiler, self).for_update_clause(select)
 
+    def _append_returning(self, text, stmt):
+        returning_cols = stmt.kwargs.get('postgres_returning', None)
+        if returning_cols:
+            def flatten_columnlist(collist):
+                for c in collist:
+                    if isinstance(c, expression.Selectable):
+                        for co in c.columns:
+                            yield co
+                    else:
+                        yield c
+            columns = [self.process(c) for c in flatten_columnlist(returning_cols)]
+            text += ' RETURNING ' + string.join(columns, ', ')
+        
+        return text
+
+    def visit_update(self, update_stmt):
+        text = super(PGCompiler, self).visit_update(update_stmt)
+        return self._append_returning(text, update_stmt)
+
+    def visit_insert(self, insert_stmt):
+        text = super(PGCompiler, self).visit_insert(insert_stmt)
+        return self._append_returning(text, insert_stmt)
+
 class PGSchemaGenerator(compiler.SchemaGenerator):
     def get_column_specification(self, column, **kwargs):
         colspec = self.preparer.format_column(column)
diff --git a/lib/sqlalchemy/sql/expression.py b/lib/sqlalchemy/sql/expression.py
index 9a9cf65d2..6f3ee94ab 100644
--- a/lib/sqlalchemy/sql/expression.py
+++ b/lib/sqlalchemy/sql/expression.py
@@ -249,7 +249,7 @@ def subquery(alias, *args, **kwargs):
 
     return Select(*args, **kwargs).alias(alias)
 
-def insert(table, values=None, inline=False):
+def insert(table, values=None, inline=False, **kwargs):
     """Return an [sqlalchemy.sql.expression#Insert] clause element.
 
     Similar functionality is available via the ``insert()`` method on
@@ -287,9 +287,9 @@ def insert(table, values=None, inline=False):
     against the ``INSERT`` statement.
     """
 
-    return Insert(table, values, inline=inline)
+    return Insert(table, values, inline=inline, **kwargs)
 
-def update(table, whereclause=None, values=None, inline=False):
+def update(table, whereclause=None, values=None, inline=False, **kwargs):
     """Return an [sqlalchemy.sql.expression#Update] clause element.
 
     Similar functionality is available via the ``update()`` method on
@@ -332,7 +332,7 @@ def update(table, whereclause=None, values=None, inline=False):
     against the ``UPDATE`` statement.
     """
 
-    return Update(table, whereclause=whereclause, values=values, inline=inline)
+    return Update(table, whereclause=whereclause, values=values, inline=inline, **kwargs)
 
 def delete(table, whereclause = None, **kwargs):
     """Return a [sqlalchemy.sql.expression#Delete] clause element.
@@ -2699,11 +2699,11 @@ class TableClause(FromClause):
     def select(self, whereclause = None, **params):
         return select([self], whereclause, **params)
 
-    def insert(self, values=None, inline=False):
-        return insert(self, values=values, inline=inline)
+    def insert(self, values=None, inline=False, **kwargs):
+        return insert(self, values=values, inline=inline, **kwargs)
 
-    def update(self, whereclause=None, values=None, inline=False):
-        return update(self, whereclause=whereclause, values=values, inline=inline)
+    def update(self, whereclause=None, values=None, inline=False, **kwargs):
+        return update(self, whereclause=whereclause, values=values, inline=inline, **kwargs)
 
     def delete(self, whereclause=None):
         return delete(self, whereclause)
@@ -3356,12 +3356,14 @@ class _UpdateBase(ClauseElement):
         return self.table.bind
 
 class Insert(_UpdateBase):
-    def __init__(self, table, values=None, inline=False):
+    def __init__(self, table, values=None, inline=False, **kwargs):
         self.table = table
         self.select = None
         self.inline=inline
         self.parameters = self._process_colparams(values)
 
+        self.kwargs = kwargs
+
     def get_children(self, **kwargs):
         if self.select is not None:
             return self.select,
@@ -3383,12 +3385,14 @@ class Insert(_UpdateBase):
         return u
 
 class Update(_UpdateBase):
-    def __init__(self, table, whereclause, values=None, inline=False):
+    def __init__(self, table, whereclause, values=None, inline=False, **kwargs):
         self.table = table
         self._whereclause = whereclause
         self.inline = inline
         self.parameters = self._process_colparams(values)
 
+        self.kwargs = kwargs
+
     def get_children(self, **kwargs):
         if self._whereclause is not None:
             return self._whereclause,