首页 > 解决方案 > 如何将 INSERT INTO ... RETURNING 值分配给变量

问题描述

我正在使用 SQL iPython Magic 扩展(postrgresql)在 Python 上的 Jupyter Notebook 中做一个项目。我有一个库存表

%%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)

我在 for 循环中创建数据以填充我的库存和钻石表。

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_id 将值插入到钻石表中。

我知道我需要在将值插入到库存表后立即捕获 inventory_id。我担心这部分代码

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 分配给变量并使用它来填充 diamonds 表吗?

标签: pythonsqlpostgresql

解决方案


@Srinika 是正确的。这可以通过 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 子句),该表达式插入到返回生成的 inventory_id 的库存中。主要部分选择其他参数以及返回的 inventory_id 并插入到 diamonds 中。没有返回,由函数头中的“returns void”表示。

现在我的 Python 很弱,但我相信单语句格式是有效的。我认为会导致类似的结果,尽管我完全不确定缩进,但请记住整个事情是 1 条语句。

 %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; 

我建议进行一个建模更改:从库存中删除价格。如果我有超过 1 颗钻石,该列应该包含什么。或者发生的事情是删除一个(或可能出售)。这是一个可导出的值,可以根据需要进行计算。


推荐阅读