diff options
author | Ivan Kotelnikov <ivan.kotelnikov@bitfury.com> | 2018-11-21 16:29:34 +0200 |
---|---|---|
committer | Ivan Kotelnikov <ivan.kotelnikov@bitfury.com> | 2018-11-21 16:38:13 +0200 |
commit | 9a8d7d5a6ad5f24e1cbc529e98dd61ad9424b9ed (patch) | |
tree | 236b61526527daa564391a8d056d68989156270e /lib | |
parent | 483901ea7bb796141c4d848f3d8b930d8ab8b32e (diff) | |
download | psycopg2-9a8d7d5a6ad5f24e1cbc529e98dd61ad9424b9ed.tar.gz |
add RETURNING clause support in execute_values function
Diffstat (limited to 'lib')
-rw-r--r-- | lib/extras.py | 11 |
1 files changed, 10 insertions, 1 deletions
diff --git a/lib/extras.py b/lib/extras.py index 0764edf..3e113b0 100644 --- a/lib/extras.py +++ b/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_result=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,9 @@ 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_result: flag indicating that results of query execution should + be returned. Useful for queries with `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 +1268,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 = [] 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 +1278,11 @@ 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: + result.extend(cur.fetchall()) + + if fetch_result: + return result def _split_sql(sql): |