使用psycopg2获取插入的多行数据的id

Question

我想使用psycopg2来一次性插入多行数据，并使用单个查询返回所有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，而不仅仅是一个？

- Jian

---

有趣的问题。请问您使用的是哪个 psycopg2 版本以及底层 PostgreSQL 版本呢？ - Craig Ringer

psycopg2 2.4.5 (dt dec pq3 ext) 和 PostgreSQL 9.2.4 - Jian

Answer 1

您不应该能够从executemany中获取结果：

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

根据psycopg2文档。

最好循环在事务内执行单个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进行测试。

- Craig Ringer

---

谢谢！我完全不知道RETURNING的顺序是不保证的。不幸的是，我无法唯一匹配返回的id。在这种情况下如何继续进行（除了插入一个虚拟的唯一字段）？ - Jian

@Jian 如果您需要将生成的ID与插入的元组关联起来，您需要逐个插入它们，或者使用returning *并匹配整个插入的元组。(Pg目前按照values子句中出现的输入顺序返回ID，但它可能会在未来停止这样做，而SQL规范不要求这样做)。 - Craig Ringer

@CraigRinger，您能指出PostgreSQL文档中是否有说明RETURNING返回的行是否按照给定的VALUES顺序排列吗？如果不能保证这一点，那将会是一个真正的痛点。 - Ben Hoyt

这就是关键所在 - 它没有明确保证任何地方，并且在 SQL 中，除非特别排序，否则一切都是无序集合。RETURNING 顺序可能会改变的原因可能是引入支持索引排序表之类的东西。不过最好在 pgsql-general 上提出这个问题 - 可能 SQL 规范本身提供了这样的保证，在这种情况下，Pg 必须遵守它。现在，它似乎只是因为实现将按照输入时呈现的顺序消耗和发出行而工作。 - Craig Ringer

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

- Gregor Zeitlinger

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

- Clodoaldo Neto

---

不错！我一直在想如何接受动态生成的数据，而不必使用Python字符串格式化，这正是所需。但是，如果您事先不知道字段名称，因为它们也是动态生成的呢？是否有一个干净的psycopg2 API可以格式化它们？ - Jian