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.

all ```CTE```s in Postgresql always work with original data snapshot, and the real data change is done at the end of execution. so if you want the second CTE to see the results of the first, then the first must be with RETURNING and the second must select from first. — Игорь Тыра, Aug 29 '20 at 16:43

score 1 · Answer 1 · answered Aug 29 '20 at 13:13

I don't think such construct can work in Postgres; even in MySQL, it wouldn't be possible - or at least safe - to use variables this way.

I think I understand that you want to keep track the total stock available before the update is executed. It might be simpler to just use two different queries for this:

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

If you want to avoid race conditions, you can wrap those in a transaction, or write it as a single statement:

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

I was trying a for-in-select-loop, that it is similar to your solution. — coterobarros, Aug 29 '20 at 13:16
However, the idea is to keep track of the stock **moved** from available to locked, not of all the stock_available at the end, that should be 0. — coterobarros, Aug 29 '20 at 13:19

bfris · Accepted Answer · 2020-08-31T16:05:43.733

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|

Interesting aproach @bfris. In my mind creating a temporal table to store a number sounds costly but surely is an idea from ancient times, as I have read elsewhere that temporal tables in PostgreSQL are cheap in computing terms. Thanks! — coterobarros, Aug 29 '20 at 16:42
Temporary tables wasn't my first choice. But they can be VERY useful for some problems. For non privileged users, the SQL Only method might be one of the few ways to solve this. — bfris, Aug 29 '20 at 16:55

score 0 · Answer 3 · answered Aug 29 '20 at 13:24

This is a dirty solution that executes as many updates as lines in table stocks. It works but this is the kind of solution I was trying to avoid.

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|

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

3 Answers3

Using Only SQL

Using Stored Procedure