获取插入到 psycopg2 中的多行的 id

Question

我想将 psycopg2 用于INSERT多行，然后id使用单个查询返回所有s（按顺序）。这就是 PostgreSQLRETURNING扩展的设计目的，使用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,)]

我如何让它正确返回所有ids 而不是一个？

Answer 1

您不应该能够从executemany以下位置获得结果：

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

根据psycopg2 文档。

您最好insert在事务中循环单个或使用 multi-valued 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 2

诀窍是使用 mogrify。它使用单个 execute 和 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()]