получение идентификаторов нескольких строк, вставленных в psycopg2

Я хотел бы использовать 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,)]

Как мне заставить его правильно возвращать все id вместо одного?


person Jian    schedule 07.02.2014    source источник
comment
Интересная проблема. версия psycopg2 и базовая клиентская версия PostgreSQL?   -  person Craig Ringer    schedule 07.02.2014
comment
psycopg2 2.4.5 (расширение dt dec pq3) и PostgreSQL 9.2.4   -  person Jian    schedule 07.02.2014


Ответы (3)


Вы не должны получать результаты от executemany:

Эта функция в основном полезна для команд, которые обновляют базу данных: любой набор результатов, возвращаемый запросом, отбрасывается.

Согласно документации psycopg2.

Вам будет лучше зацикливаться на одном insert в транзакции или использовать многозначное insert... returning, хотя в последнем случае вы должны быть осторожны, чтобы сопоставлять возвращаемые идентификаторы с использованием другого входного значения, вы не можете просто принять порядок возвращаемые идентификаторы совпадают со входным списком 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.

person Craig Ringer    schedule 07.02.2014
comment
Спасибо! Я совершенно не знал, что заказ RETURNING не является гарантированным. К сожалению, у меня нет возможности однозначно сопоставить возвращенные ids. Любые идеи о том, как действовать в этом случае (кроме простой вставки фиктивного уникального поля)? - person Jian; 08.02.2014
comment
@Jian Если вам нужно связать возвращенный сгенерированный идентификатор со вставленным кортежем, вам нужно просто вставить их один за другим или использовать returning * и сопоставить весь вставленный кортеж. (В настоящее время Pg возвращает идентификаторы в том порядке, в котором входные данные появляются в предложении values, но это может прекратиться в будущем, и спецификация SQL не требует этого). - person Craig Ringer; 08.02.2014
comment
@CraigRinger, можете ли вы указать на документы PostgreSQL, где говорится, соответствуют ли строки, возвращаемые RETURNING, указанным вами ЗНАЧЕНИЯМ? Было бы настоящей болью, если бы это не было гарантировано. - person Ben Hoyt; 31.03.2014
comment
В том-то и дело - это нигде явно не гарантируется, а в SQL все является неупорядоченным набором, если специально не заказано. Причинами, по которым порядок RETURNING может измениться, могут быть такие вещи, как введение поддержки таблиц, упорядоченных по индексу. Тем не менее, было бы неплохо поднять вопрос о pgsql-general - возможно, спецификация SQL сама по себе предоставляет такую ​​​​гарантию, и в этом случае Pg должен ее соблюдать. Прямо сейчас это работает только потому, что реализация будет потреблять и выдавать строки в том порядке, в котором они представлены в качестве входных данных. - person Craig Ringer; 01.04.2014

Хитрость заключается в использовании mogrify. Он использует одно выполнение и идентификатор, поэтому в любом случае быстрее, чем 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()]
person Gregor Zeitlinger    schedule 16.03.2015

Передайте динамически сгенерированные данные в виде массива кортежей и удалите их.

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,)]
person Clodoaldo Neto    schedule 07.02.2014
comment
Хороший! Мне было интересно, как принимать динамически сгенерированные данные, не прибегая к форматированию строк Python, и это делает именно это. Но что, если вы не знаете имена полей заранее, потому что они также генерируются динамически? Есть ли чистый API psycopg2 для их форматирования? - person Jian; 08.02.2014