I'd like to use psycopg2 to INSERT multiple rows and then return all the ids (in order) using a single query. This is what PostgreSQL's RETURNING extension is designed for, and it seems to work fine using cursor.execute:
cursor.execute(
"INSERT INTO my_table (field_1, field_2) "
"VALUES (0, 0), (0, 0) RETURNING id;"
)
print cursor.fetchall()
[(1,), (2,)]
Now, in order to pass in dynamically-generated data, it seems like cursor.executemany is the way to go:
data = [(0, 0), (0, 0)]
cursor.executemany(
"INSERT INTO my_table (field_1, field_2) "
"VALUES (%s, %s) RETURNING id;",
data
)
However, in that case, cursor.fetchall() produces the following:
[(4,), (None,)]
How do I get it to correctly return all the ids instead of just the one?
You're not supposed to be able to get results from executemany:
The function is mostly useful for commands that update the database: any result set returned by the query is discarded.
Per the psycopg2 docs.
You'll be better off looping over a single insert within a transaction, or using a multi-valued insert... returning, though in the latter case you must be careful to match returned IDs using another input value, you can't just assume the order of returned IDs is the same as the input VALUES list.
When I run your test locally, it simply fails:
>>> import psycopg2
>>> conn = psycopg2.connect("dbname=regress")
>>> curs = conn.cursor()
>>> curs.execute("create table my_table(id serial primary key, field_1 integer, field_2 integer);")
>>> data = [(0, 0), (0, 0)]
>>> curs.executemany(
... "INSERT INTO my_table (field_1, field_2) "
... "VALUES (%s, %s) RETURNING id;",
... data
... )
>>>
>>> curs.fetchall()
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
psycopg2.ProgrammingError: no results to fetch
Tested with psycopg2 2.5.1.
Pass the dynamically-generated data as an array of tuples and unnest it
import psycopg2
insert = """
insert into my_table (field_1, field_2)
select field_1, field_2
from unnest(%s) s(field_1 int, field_2 int)
returning id
;"""
data = [(0,0),(1,1),(2,2)]
conn = psycopg2.connect("host=localhost4 port=5432 dbname=cpn")
cursor = conn.cursor()
cursor.execute(insert, (data,))
print cursor.fetchall()
conn.commit()
conn.close()
Prints
[(1,), (2,), (3,)]
The trick is to use mogrify. It uses a single execute and id therefore faster than executemany anyways:
def insert_many(self, table: str, id_column: str, values: list):
if not values:
return []
keys = values[0].keys()
query = cursor.mogrify("INSERT INTO {} ({}) VALUES {} RETURNING {}".format(
table,
', '.join(keys),
', '.join(['%s'] * len(values)),
id_column
), [tuple(v.values()) for v in values])
conn = psycopg2.connect("host=localhost4 port=5432 dbname=cpn")
cursor = conn.cursor()
cursor.execute(query)
return [t[0] for t in (cursor.fetchall()]
