How to assign a local variable in an update sentence in PostgreSQL

Question

I was trying to assign a local variable (well, in fact two) using the SET clause of an update sentence running on several rows. Ok, I am doing this 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;
$$;

And unluckily this is not the way PostgreSQL expects to do such a thing.

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

This is only a simplistic example to illustrate the real problem of counting/summing-up the number of things updated in an update sentence. I am sure there can be tricks or ways around for this specific example, but I am interested in the general solution for situations like this, where an accumulator must be calculated.

Any idea?

---

EDIT

Following @GMB suggestion, I chain 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;
$$;

This should work but the resulting value indicates that both selects are executed on the preliminary values of table stocks, as the difference is 0.

select lock_all3();

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

However, cte2 is executed, as the final situation indicates that all available stocks were locked.

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|

There still must be something wrong in this approximation.

Answer 1

I think the trick is to calculate the total before the update.

Using Only 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|

Using Stored Procedure

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|

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|