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

Question

我想使用 psycopg2 到 INSERT 多行，然后使用单个查询返回所有 ids（按顺序）。这就是 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,)]

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

Answer 1

您不应该能够从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 2

诀窍是使用 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()]

Answer 3

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

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