Merge remote-tracking branch 'eternalflow/execute-values-returning-clause-support'

NEWS

@@ -18,6 +18,8 @@ New features:
   structure (:ticket:`#782`).
 - `~psycopg2.sql.Identifier` can represent qualified names in SQL composition
   (:ticket:`#732`).
+- Added *fetch* parameter to `~psycopg2.extras.execute_values()` function
+  (:ticket:`813`).
 - Fixed adaptation of numeric subclasses such as `IntEnum` (:ticket:`591`).
 - `!str()` on `~psycopg2.extras.Range` produces a human-readable representation
   (:ticket:`#773`).

doc/src/extras.rst

@@ -1023,6 +1023,8 @@ parameters.  By reducing the number of server roundtrips the performance can be
 .. autofunction:: execute_values
 
     .. versionadded:: 2.7
+    .. versionchanged:: 2.8
+        added the *fetch* parameter.
 
 
 .. index::

lib/extras.py

@@ -1198,7 +1198,7 @@ def execute_batch(cur, sql, argslist, page_size=100):
         cur.execute(b";".join(sqls))
 
 
-def execute_values(cur, sql, argslist, template=None, page_size=100):
+def execute_values(cur, sql, argslist, template=None, page_size=100, fetch=False):
     '''Execute a statement using :sql:`VALUES` with a sequence of parameters.
 
     :param cur: the cursor to use to execute the query.
@@ -1229,6 +1229,10 @@ def execute_values(cur, sql, argslist, template=None, page_size=100):
         statement. If there are more items the function will execute more than
         one statement.
 
+    :param fetch: if `!True` return the query results into a list (like in a
+        `~cursor.fetchall()`).  Useful for queries with :sql:`RETURNING`
+        clause.
+
     .. __: https://www.postgresql.org/docs/current/static/queries-values.html
 
     After the execution of the function the `cursor.rowcount` property will
@@ -1265,6 +1269,7 @@ def execute_values(cur, sql, argslist, template=None, page_size=100):
         sql = sql.encode(_ext.encodings[cur.connection.encoding])
     pre, post = _split_sql(sql)
 
+    result = [] if fetch else None
     for page in _paginate(argslist, page_size=page_size):
         if template is None:
             template = b'(' + b','.join([b'%s'] * len(page[0])) + b')'
@@ -1274,6 +1279,10 @@ def execute_values(cur, sql, argslist, template=None, page_size=100):
             parts.append(b',')
         parts[-1:] = post
         cur.execute(b''.join(parts))
+        if fetch:
+            result.extend(cur.fetchall())
+
+    return result
 
 
 def _split_sql(sql):

tests/test_fast_executemany.py

@@ -229,6 +229,15 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
         cur.execute("select id, data from testfast where id = 3")
         self.assertEqual(cur.fetchone(), (3, snowman))
 
+    def test_returning(self):
+        cur = self.conn.cursor()
+        result = psycopg2.extras.execute_values(cur,
+            "insert into testfast (id, val) values %s returning id",
+            ((i, i * 10) for i in range(25)),
+            page_size=10, fetch=True)
+        # result contains all returned pages
+        self.assertEqual([r[0] for r in result], list(range(25)))
+
     def test_invalid_sql(self):
         cur = self.conn.cursor()
         self.assertRaises(ValueError, psycopg2.extras.execute_values, cur,