Как присвоить переменной INSERT INTO ... RETURNING значение

Question

Я делаю проект в Jupyter Notebook на Python с расширением SQL i Python Magi c (postrgre sql). У меня есть таблица инвентаря

%%sql
   CREATE TABLE inventory (
   inventory_id SERIAL PRIMARY KEY,
   name VARCHAR(20),
   price INT,
   sku VARCHAR(10))

и таблица алмазов

%%sql
   CREATE TABLE diamonds (
   diamond_id SERIAL PRIMARY KEY,
   shape VARCHAR(20),
   carats REAL,
   color VARCHAR(20),
   clarity VARCHAR(20),
   cut VARCHAR(20),
   price INT,
   inventory_id INT)

Я создаю данные для l oop, чтобы заполнить свои таблицы инвентаря и бриллиантов.

for i in range(10):

shape = random.choice(shape_list)
color = random.choice(color_list)
clarity = random.choice(clarity_list)
cut = random.choice(cut_list)
carats = round(random.uniform(0.1, 14),2)
if carats < 0.3:
    price = random.randint(200, 600)
elif carats < 0.5:
    price = random.randint(600, 1000)
elif carats < 0.7:
    price = random.randint(1000, 1800)
elif carats < 1.5:
    price = random.randint(1800, 5000)
elif carats < 3:
    price = random.randint(5000, 17000)
elif carats < 5:
    price = random.randint(17000, 25000)
elif carats < 8:
    price = random.randint(30000, 300000)
elif carats <= 14:
    price = random.randint(300000, 1000000)
id = 'dmd_'+str(i)
name=str(carats)+' '+'ct'+' '+'Diamond'
sku = randomString()
new_id = %sql INSERT INTO inventory (name, price, sku) VALUES (:name, :price, :sku) RETURNING inventory_id
%sql INSERT INTO diamonds (shape, carats, color, clarity, cut, price, inventory_id) VALUES (:shape, :carats, :color, :clarity, :cut, :price, :new_id)

Мне нужно получить inventory_id после вставки значений в таблицу inventory и последующего использования этих значений inventory_id вставки в таблицу diamonds.

Я понимаю, что мне нужно поймать inventory_id сразу после вставки значений в таблицу inventory. Меня беспокоит эта часть кода

new_id = %sql INSERT INTO inventory (name, price, sku) VALUES (:name, :price, :sku) RETURNING inventory_id
%sql INSERT INTO diamonds (shape, carats, color, clarity, cut, price, inventory_id) VALUES (:shape, :carats, :color, :clarity, :cut, :price, :new_id)

В этой части я получил ошибку «(psycopg2.ProgrammingError) не могу адаптировать тип RowProxy»

Есть ли другие способы присвоить inventory_id переменной и использовать ее для заполнения таблицы бриллиантов?

Answer 1

@ Сриника правильно. Это можно сделать с помощью Postgres хранимых процедур. Но можно сделать еще один шаг, поскольку это может быть один оператор SQL или функция, состоящая только из этого одного оператора. Эта функция будет выглядеть так:

create or replace function create_inventory_diamonds
                         ( name_in       varchar
                         , inv_price_in  integer
                         , sku_in        integer
                         , shape_in      varchar
                         , carats_in     real 
                         , color_in      varchar
                         , clarity_in    varchar 
                         , cut_in        varchar
                         , price_in      integer
                         )
  returns void
  language sql 
as $$ 
with inv as 
     ( insert into inventory(name, price, sku)
          values (name_in, inv_price_in, sku_in) 
        returning inventory_id
     )
insert into diamonds 
          ( shape
          , carats 
          , color 
          , clarity  
          , cut  
          , price 
          , inventory_id
          )
    select  shape_in
          , carats_in 
          , color_in  
          , clarity_in   
          , cut_in   
          , price_in  
          , inv.inventory_id  --reference back to CTE
     from inv;

$$;          

Она создает общее табличное выражение (CTE, with clause), которое вставляется в инвентарь, возвращая сгенерированный inventory_id. Основная часть выбирает другие параметры вместе с возвращенным inventory_id и вставляет их в ромбы. Возврата нет, на что указывает «возвращает пустоту» в заголовке функции.

Теперь мой Python довольно слабый, но я считаю, что формат единственного оператора будет допустимым. Я думаю, что это приведет к чему-то вроде, хотя я абсолютно не уверен в отступе, помните, что все это одно утверждение.

 %sql  with inv as ( insert into inventory(name, price, sku) VALUES (:name, :price, :sku)   
       insert into diamonds ( shape, carats, color, clarity,cut, price, inventory_id)
       select  :shape, :carats, :color, :clarity, :cut, :price, inv.inventory_id; 

Есть изменение модели, которое я бы посоветовал: удалить цену из инвентаря. Что должен содержать этот столбец, если у меня более одного алмаза. Или то, что происходит, удаляют (или, возможно, продают). Это производная стоимость, которая может быть рассчитана по мере необходимости.