Question

我想将psycopg2用于INSERT多行，然后使用单个查询返回所有id（按顺序）。这就是PostgreSQL的RETURNING扩展程序的设计目标，它似乎可以正常使用cursor.execute：

cursor.execute(
    "INSERT INTO my_table (field_1, field_2) "
    "VALUES (0, 0), (0, 0) RETURNING id;"
)
print cursor.fetchall()

[(1,), (2,)]

现在，为了传递动态生成的数据，似乎可以使用cursor.executemany：

data = [(0, 0), (0, 0)]

cursor.executemany(
    "INSERT INTO my_table (field_1, field_2) "
    "VALUES (%s, %s) RETURNING id;",
    data
)

但是，在这种情况下，cursor.fetchall()会产生以下内容：

[(4,), (None,)]

如何让它正确返回所有id而不仅仅是那个？

Answer 1

将动态生成的数据作为元组数组传递并取消它

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()

打印

[(1,), (2,), (3,)]

Answer 2

您不应该从executemany获得结果：

该函数主要用于更新数据库的命令：查询返回的任何结果集都将被丢弃。

每the psycopg2 docs。

最好在事务中循环一个insert，或使用多值insert... returning，但在后一种情况下，您必须小心使用其他输入匹配返回的ID值，您不能只假设返回的ID的顺序与输入VALUES列表相同。

当我在本地运行测试时，它只是失败了：

>>> 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

使用psycopg2 2.5.1进行测试。

Answer 3

诀窍是使用mogrify。它使用单个执行和id因此比executemany更快：

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()]

获取psycopg2中插入的多行ID

3 个答案: