Merge lp:~ack/storm/update-returning into lp:storm

[Alberto Donato]

This branch adds support for Returning(Update(...)) expressions, with an optional list of columns to return.

[Free Ekanayaka (free.ekanayaka)]

Looks good, +1!

[1]

Please update the NEWS file, mentioning this new feature.

[2]

- """Appends the "RETURNING <primary_columns>" suffix to an INSERT.
+ """Appends the "RETURNING <primary_columns>" suffix to an INSERT or UPDATE.

This is slightly inaccurate now, as you can have arbitrary columns. I'd change it to:

    """Appends the "RETURNING <columns>" suffix to an INSERT or UPDATE.

    @param expression: ...
    @param columns: The columns to return, if C{None} then C{expression.primary_columns}
        will be used.

[3]

+ def __init__(self, expression, columns=None):

In this particular case I'd s/expression/expr/, as the class name is Expr, and
the current convention in the code is to use the name "expr" for instances of Expr.

review: Approve

[Jamu Kakar (jkakar)]

Nice work, +1! I have the same comments as Free, please fix them
before merging.

review: Approve

[lp:~ack/storm/update-returning]

436. By Alberto Donato on 2012-03-10

Addressed review comments.

[Alberto Donato (ack)]

> Looks good, +1!
>
> [1]
>
> Please update the NEWS file, mentioning this new feature.
>
> [2]
>
> - """Appends the "RETURNING <primary_columns>" suffix to an INSERT.
> + """Appends the "RETURNING <primary_columns>" suffix to an INSERT or
> UPDATE.
>
> This is slightly inaccurate now, as you can have arbitrary columns. I'd change
> it to:
>
> """Appends the "RETURNING <columns>" suffix to an INSERT or UPDATE.
>
> @param expression: ...
> @param columns: The columns to return, if C{None} then
> C{expression.primary_columns}
> will be used.
>
> [3]
>
> + def __init__(self, expression, columns=None):
>
> In this particular case I'd s/expression/expr/, as the class name is Expr, and
> the current convention in the code is to use the name "expr" for instances of
> Expr.

Addressed all points, thanks!

[lp:~ack/storm/update-returning]

437. By Alberto Donato on 2012-03-10

Update NEWS file.

=== modified file 'NEWS'
--- NEWS	2012-03-08 10:22:31 +0000
+++ NEWS	2012-03-10 15:24:18 +0000
@@ -26,6 +26,10 @@
 - Insert expressions now support multi-row and subquery INSERT
   statements.
 
+- Support in the postgres backend to use the RETURNING extension for UPDATE
+  statement, optionally specifying the columns to return.
+  (PostgreSQL >= 8.2 only)
+
 Bug fixes
 ---------
 

=== modified file 'storm/databases/postgres.py'
--- storm/databases/postgres.py	2012-01-25 19:52:04 +0000
+++ storm/databases/postgres.py	2012-03-10 15:24:18 +0000
@@ -56,23 +56,28 @@
 
 
 class Returning(Expr):
-    """Appends the "RETURNING <primary_columns>" suffix to an INSERT.
+    """Appends the "RETURNING <columns>" suffix to an INSERT or UPDATE.
+
+    @param expr: an L{Insert} or L{Update} expression.
+    @param columns: The columns to return, if C{None} then
+        C{expr.primary_columns} will be used.
 
     This is only supported in PostgreSQL 8.2+.
     """
-
-    def __init__(self, insert):
-        self.insert = insert
+    def __init__(self, expr, columns=None):
+        self.expr = expr
+        self.columns = columns
 
 @compile.when(Returning)
 def compile_returning(compile, expr, state):
     state.push("context", COLUMN)
-    columns = compile(expr.insert.primary_columns, state)
+    columns = expr.columns or expr.expr.primary_columns
+    columns = compile(columns, state)
     state.pop()
     state.push("precedence", 0)
-    insert = compile(expr.insert, state)
+    expr = compile(expr.expr, state)
     state.pop()
-    return "%s RETURNING %s" % (insert, columns)
+    return "%s RETURNING %s" % (expr, columns)
 
 
 class currval(FuncExpr):

=== modified file 'storm/expr.py'
--- storm/expr.py	2012-03-08 10:09:38 +0000
+++ storm/expr.py	2012-03-10 15:24:18 +0000
@@ -749,13 +749,15 @@
 
 
 class Update(Expr):
-    __slots__ = ("map", "where", "table", "default_table")
+    __slots__ = ("map", "where", "table", "default_table", "primary_columns")
 
-    def __init__(self, map, where=Undef, table=Undef, default_table=Undef):
+    def __init__(self, map, where=Undef, table=Undef, default_table=Undef,
+                 primary_columns=Undef):
         self.map = map
         self.where = where
         self.table = table
         self.default_table = default_table
+        self.primary_columns = primary_columns
 
 @compile.when(Update)
 def compile_update(compile, update, state):

=== modified file 'tests/databases/postgres.py'
--- tests/databases/postgres.py	2012-01-25 22:03:17 +0000
+++ tests/databases/postgres.py	2012-03-10 15:24:18 +0000
@@ -29,7 +29,7 @@
 from storm.variables import ListVariable, IntVariable, Variable
 from storm.properties import Int
 from storm.exceptions import DisconnectionError
-from storm.expr import (Union, Select, Insert, Alias, SQLRaw, State,
+from storm.expr import (Union, Select, Insert, Update, Alias, SQLRaw, State,
                         Sequence, Like, Column, COLUMN)
 from storm.tracer import install_tracer, TimeoutError
 from storm.uri import URI
@@ -90,13 +90,13 @@
                                 "(id SERIAL PRIMARY KEY, b BYTEA)")
         self.connection.execute("CREATE TABLE like_case_insensitive_test "
                                 "(id SERIAL PRIMARY KEY, description TEXT)")
-        self.connection.execute("CREATE TABLE insert_returning_test "
+        self.connection.execute("CREATE TABLE returning_test "
                                 "(id1 INTEGER DEFAULT 123, "
                                 " id2 INTEGER DEFAULT 456)")
 
     def drop_tables(self):
         super(PostgresTest, self).drop_tables()
-        for table in ["like_case_insensitive_test", "insert_returning_test"]:
+        for table in ["like_case_insensitive_test", "returning_test"]:
             try:
                 self.connection.execute("DROP TABLE %s" % table)
                 self.connection.commit()
@@ -389,25 +389,32 @@
                           "thetable.thecolumn = "
                           "(SELECT currval('thetable_thecolumn_seq'))")
 
-    def test_returning(self):
-        insert = Insert({column1: elem1}, table1,
-                        primary_columns=(column2, column3))
-        self.assertEquals(compile(Returning(insert)),
-                          'INSERT INTO "table 1" (column1) VALUES (elem1) '
-                          'RETURNING column2, column3')
-
     def test_returning_column_context(self):
         column2 = TrackContext()
         insert = Insert({column1: elem1}, table1, primary_columns=column2)
         compile(Returning(insert))
         self.assertEquals(column2.context, COLUMN)
 
+    def test_returning_update(self):
+        update = Update({column1: elem1}, table=table1,
+                        primary_columns=(column2, column3))
+        self.assertEquals(compile(Returning(update)),
+                          'UPDATE "table 1" SET column1=elem1 '
+                          'RETURNING column2, column3')
+
+    def test_returning_update_with_columns(self):
+        update = Update({column1: elem1}, table=table1,
+                        primary_columns=(column2, column3))
+        self.assertEquals(compile(Returning(update, columns=[column3])),
+                          'UPDATE "table 1" SET column1=elem1 '
+                          'RETURNING column3')
+
     def test_execute_insert_returning(self):
-        if self.database._version < (8, 2):
+        if self.database._version < 80200:
             return # Can't run this test with old PostgreSQL versions.
 
-        column1 = Column("id1", "insert_returning_test")
-        column2 = Column("id2", "insert_returning_test")
+        column1 = Column("id1", "returning_test")
+        column2 = Column("id2", "returning_test")
         variable1 = IntVariable()
         variable2 = IntVariable()
         insert = Insert({}, primary_columns=(column1, column2),
@@ -420,13 +427,13 @@
         self.assertEquals(variable1.get(), 123)
         self.assertEquals(variable2.get(), 456)
 
-        result = self.connection.execute("SELECT * FROM insert_returning_test")
+        result = self.connection.execute("SELECT * FROM returning_test")
         self.assertEquals(result.get_one(), (123, 456))
 
     def test_wb_execute_insert_returning_not_used_with_old_postgres(self):
         """Shouldn't try to use RETURNING with PostgreSQL < 8.2."""
-        column1 = Column("id1", "insert_returning_test")
-        column2 = Column("id2", "insert_returning_test")
+        column1 = Column("id1", "returning_test")
+        column2 = Column("id2", "returning_test")
         variable1 = IntVariable()
         variable2 = IntVariable()
         insert = Insert({}, primary_columns=(column1, column2),
@@ -438,31 +445,44 @@
         self.assertFalse(variable1.is_defined())
         self.assertFalse(variable2.is_defined())
 
-        result = self.connection.execute("SELECT * FROM insert_returning_test")
+        result = self.connection.execute("SELECT * FROM returning_test")
         self.assertEquals(result.get_one(), (123, 456))
 
     def test_execute_insert_returning_without_columns(self):
         """Without primary_columns, the RETURNING system won't be used."""
-        column1 = Column("id1", "insert_returning_test")
+        column1 = Column("id1", "returning_test")
         variable1 = IntVariable()
         insert = Insert({column1: 123}, primary_variables=(variable1,))
         self.connection.execute(insert)
 
         self.assertFalse(variable1.is_defined())
 
-        result = self.connection.execute("SELECT * FROM insert_returning_test")
+        result = self.connection.execute("SELECT * FROM returning_test")
         self.assertEquals(result.get_one(), (123, 456))
 
     def test_execute_insert_returning_without_variables(self):
         """Without primary_variables, the RETURNING system won't be used."""
-        column1 = Column("id1", "insert_returning_test")
+        column1 = Column("id1", "returning_test")
         insert = Insert({}, primary_columns=(column1,))
         self.connection.execute(insert)
 
-        result = self.connection.execute("SELECT * FROM insert_returning_test")
+        result = self.connection.execute("SELECT * FROM returning_test")
 
         self.assertEquals(result.get_one(), (123, 456))
 
+    def test_execute_update_returning(self):
+        if self.database._version < 80200:
+            return # Can't run this test with old PostgreSQL versions.
+
+        column1 = Column("id1", "returning_test")
+        column2 = Column("id2", "returning_test")
+        self.connection.execute(
+            "INSERT INTO returning_test VALUES (1, 2)")
+        update = Update({"id2": 3}, column1 == 1,
+                        primary_columns=(column1, column2))
+        result = self.connection.execute(Returning(update))
+        self.assertEquals(result.get_one(), (1, 3))
+
     def test_isolation_autocommit(self):
         database = create_database(
             os.environ["STORM_POSTGRES_URI"] + "?isolation=autocommit")

=== modified file 'tests/store/postgres.py'
--- tests/store/postgres.py	2011-02-14 12:17:54 +0000
+++ tests/store/postgres.py	2012-03-10 15:24:18 +0000
@@ -183,7 +183,7 @@
         Ensure that the currval()-based identity retrieval continues
         to work, even if we're currently running on a 8.2+ database.
         """
-        self.database._version = (8, 1, 9)
+        self.database._version = 80109
         foo1 = self.store.add(Foo())
         self.store.flush()
         foo2 = self.store.add(Foo())