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

2 files changed, 73 insertions, 1 deletions

diff --git a/test/sql/query.py b/test/sql/query.py
index a519dd974..b4afbbade 100644
--- a/test/sql/query.py
+++ b/test/sql/query.py
@@ -563,7 +563,51 @@ class QueryTest(PersistTest):
         s = users.select(users.c.user_name.in_() == None)
         r = s.execute().fetchall()
         assert len(r) == 1
-        
+
+    @testing.supported('postgres')
+    def test_update_returning(self):
+        meta = MetaData(testbase.db)
+        table = Table('tables', meta, 
+            Column('id', Integer, primary_key=True),
+            Column('persons', Integer),
+            Column('full', Boolean)
+        )
+        table.create()
+        try:
+            table.insert().execute([{'persons': 5, 'full': False}, {'persons': 3, 'full': False}])
+            
+            result = table.update(table.c.persons > 4, dict(full=True), postgres_returning=[table.c.id]).execute()
+            self.assertEqual(result.fetchall(), [(1,)])
+            
+            result2 = select([table.c.id, table.c.full]).order_by(table.c.id).execute()
+            self.assertEqual(result2.fetchall(), [(1,True),(2,False)])
+        finally:
+            table.drop()
+
+    @testing.supported('postgres')
+    def test_insert_returning(self):
+        meta = MetaData(testbase.db)
+        table = Table('tables', meta, 
+            Column('id', Integer, primary_key=True),
+            Column('persons', Integer),
+            Column('full', Boolean)
+        )
+        table.create()
+        try:
+            result = table.insert(postgres_returning=[table.c.id]).execute({'persons': 1, 'full': False})
+            
+            self.assertEqual(result.fetchall(), [(1,)])
+            
+            # Multiple inserts only return the last row
+            result2 = table.insert(postgres_returning=[table]).execute(
+                 [{'persons': 2, 'full': False}, {'persons': 3, 'full': True}])
+             
+            self.assertEqual(result2.fetchall(), [(3,3,True)])
+            
+            result3 = table.insert(postgres_returning=[(table.c.id*2).label('double_id')]).execute({'persons': 4, 'full': False})
+            self.assertEqual([dict(row) for row in result3], [{'double_id':8}])
+        finally:
+            table.drop()
 
 class CompoundTest(PersistTest):
     """test compound statements like UNION, INTERSECT, particularly their ability to nest on
diff --git a/test/sql/select.py b/test/sql/select.py
index 1114f1735..4cdac97d8 100644
--- a/test/sql/select.py
+++ b/test/sql/select.py
@@ -1184,7 +1184,35 @@ class CRUDTest(SQLCompileTest):
         s = select([table2.c.othername], table2.c.otherid == table1.c.myid)
         u = table1.update(table1.c.name==s)
         self.assert_compile(u, "UPDATE mytable SET myid=:myid, name=:name, description=:description WHERE mytable.name = (SELECT myothertable.othername FROM myothertable WHERE myothertable.otherid = mytable.myid)")
+    
+    @testing.supported('postgres')
+    def testupdatereturning(self):
+        dialect = postgres.dialect()
+        
+        u = update(table1, values=dict(name='foo'), postgres_returning=[table1.c.myid, table1.c.name])
+        self.assert_compile(u, "UPDATE mytable SET name=%(name)s RETURNING mytable.myid, mytable.name", dialect=dialect)
+        
+        u = update(table1, values=dict(name='foo'), postgres_returning=[table1])
+        self.assert_compile(u, "UPDATE mytable SET name=%(name)s "\
+            "RETURNING mytable.myid, mytable.name, mytable.description", dialect=dialect)
+        
+        u = update(table1, values=dict(name='foo'), postgres_returning=[func.length(table1.c.name)])
+        self.assert_compile(u, "UPDATE mytable SET name=%(name)s RETURNING length(mytable.name)", dialect=dialect)
         
+    @testing.supported('postgres')
+    def testinsertreturning(self):
+        dialect = postgres.dialect()
+        
+        i = insert(table1, values=dict(name='foo'), postgres_returning=[table1.c.myid, table1.c.name])
+        self.assert_compile(i, "INSERT INTO mytable (name) VALUES (%(name)s) RETURNING mytable.myid, mytable.name", dialect=dialect)
+        
+        i = insert(table1, values=dict(name='foo'), postgres_returning=[table1])
+        self.assert_compile(i, "INSERT INTO mytable (name) VALUES (%(name)s) "\
+            "RETURNING mytable.myid, mytable.name, mytable.description", dialect=dialect)
+        
+        i = insert(table1, values=dict(name='foo'), postgres_returning=[func.length(table1.c.name)])
+        self.assert_compile(i, "INSERT INTO mytable (name) VALUES (%(name)s) RETURNING length(mytable.name)", dialect=dialect)
+    
     def testdelete(self):
         self.assert_compile(delete(table1, table1.c.myid == 7), "DELETE FROM mytable WHERE mytable.myid = :mytable_myid")