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

Question

我试图使用 SET 分配一个局部变量(好吧，实际上是两个) update的条款句子在几行上运行。好的，我正在做这个 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;
$$;

select * from 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|