如何在 PostgreSQL 的更新语句中分配局部变量

Question

我试图使用在几行上运行的update 语句的SET 子句分配一个局部变量（实际上是两个）。好的，我正在做这个 ala MySQL。

drop table if exists stocks cascade;

create table stocks (
  id              serial,
  stock_available int,
  stock_locked    int
);

insert into stocks(stock_available, stock_locked) values 
(150, 10),
(150, 20),
(150, 0),
(100, 0),
(100, 100),
(100, 30),
(100, 0),
(100, 50),
(100, 0);

create or replace function lock_all ()
returns int
language plpgsql as $$
declare
  _amount int;
  _total int;
begin
  -- initialize accumulator
  _total = 0;

  -- update all the stocks table rows
  update stocks
  set    _amount = stock_available,
         stock_locked = stock_locked + _amount,
         _total = _total + _amount;
  
  -- returns the units locked 
  return _total;
end;
$$;

不幸的是，这不是 PostgreSQL 期望的方式。

 SQL Error [42703]: ERROR: column "_amount" of relation "stocks" does not exist
  Where: PL/pgSQL function lock_all() line 10 at SQL statement

这只是一个简单的例子来说明计算/总结update 句子中更新的事物数量的真正问题。我确信这个特定示例可能有技巧或方法，但我对这种必须计算累加器的情况的通用解决方案感兴趣。

有什么想法吗？

---

编辑

根据@GMB 的建议，我链接了 3 个 ctes

create or replace function lock_all3 ()
returns int
language sql as $$
  with 
    cte1 as (
      select 
        sum(stock_locked)::int as initially_locked 
      from 
        stocks
    ), 
    cte2 as (
      update 
        stocks 
      set 
        stock_locked = stock_locked + stock_available,
        stock_available = 0
      returning 
        0 as dummy
    ),
    cte3 as (
      select 
        sum(stock_locked)::int as finally_locked 
      from 
        stocks
    )
  select 
    (cte3.finally_locked - initially_locked - dummy) 
  from 
    cte1, cte2, cte3;
$$;

这应该可行，但结果值表明两个选择都是在表 stocks 的初始值上执行的，因为差异为 0。

select lock_all3();

lock_all3|
---------|
        0|

但是，cte2 被执行，因为最终情况表明所有可用的股票都被锁定了。

select * from stocks;

id|stock_available|stock_locked|
--|---------------|------------|
 1|              0|         160|
 2|              0|         170|
 3|              0|         150|
 4|              0|         100|
 5|              0|         200|
 6|              0|         130|
 7|              0|         100|
 8|              0|         150|
 9|              0|         100|

这个近似值肯定有问题。

Answer 1

我认为诀窍是计算更新之前的总数。

仅使用 SQL

  DROP TABLE x;
SELECT sum(stock_available) as total_moved
  INTO TEMP TABLE x
  FROM stocks as total_moved; 
UPDATE stocks 
   SET stock_locked = stock_available + stock_locked,
       stock_available = 0;
SELECT * from x;

total_moved|
-----------|
       1050|

使用存储过程

create or replace function lock_all ()
returns int
language plpgsql as $$
declare
  _total int;
begin
    --calculate total before update
    SELECT sum(stock_available)
      INTO _total
      FROM stocks;

    UPDATE stocks
       SET stock_locked = stock_locked + stock_available,
           stock_available = 0;

    return _total;
end;
$$;

从 lock_all 中选择 *；

lock_all|
--------|
    1050|

从股票中选择*；

id|stock_available|stock_locked|
--|---------------|------------|
 1|              0|         160|
 2|              0|         170|
 3|              0|         150|
 4|              0|         100|
 5|              0|         200|
 6|              0|         130|
 7|              0|         100|
 8|              0|         150|
 9|              0|         100|

Answer 2

我不认为这样的结构可以在 Postgres 中工作；即使在 MySQL 中，也不可能——或者至少是安全的——以这种方式使用变量。

我认为我了解您希望在执行更新之前跟踪可用的总库存。为此只使用两个不同的查询可能会更简单：

select sum(stock_available) total from stocks returning total into _total;
update stocks set stock_locked = stock_locked + stock_available;

如果您想避免竞争条件，可以将它们包装在事务中，或将其编写为单个语句：

with cte as (update stocks set stock_locked = stock_locked + stock_available)
select sum(stock_available) total from stocks returning total into _total;

Answer 3

这是一个脏解决方案，它执行的updates 与表stocks 中的行数一样多。它有效，但这是我试图避免的那种解决方案。

create or replace function lock_all ()
returns int
language plpgsql as $$
declare
  _amount int;
  _total int;
 r record;
begin
  -- initialize counter
  _total = 0;

    -- select stocks rows
    for r in (
        select * from stocks
    ) 
    loop
        _amount = r.stock_available;
    
      -- update the stock_fulfilled column in of_user_order_line_supply
      update stocks
      set        stock_locked = stock_locked + _amount
      where  id = r.id;
     
     _total = _total + _amount;
    end loop;

return _total;
end;
$$;

select lock_all();
select * from stocks;

lock_all|
--------|
    1050|