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

Question

Я хотел бы использовать psycopg2 для INSERT несколько строк, а затем вернуть все id s (в порядке) с использованием одного запроса. Это то, что RETURNING расширение в PostgreSQL предназначен для, и это, кажется, работает нормально, используя 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,)] 

Как я могу получить его COR прямо вернуть все id s вместо одного?

Answer 1

Вы не должны быть в состоянии получить результаты от executemany:

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

За the psycopg2 docs.

Вам будет лучше обходить цикл по одному 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.

Answer 2

Пропускают динамически сгенерированные данные в виде массива кортежей и unnest него

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 3

Хитрость заключается в том, чтобы использовать mogrify. Он использует один запуск и идентификатор, следовательно, быстрее, чем в любом случае:

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