SuperTron
SuperTron

Reputation: 4233

PostgreSQL Update and return

Let's say I have a table called t in Postgres:

   id   | group_name | state
-----------------------------
   1    |   group1   |   0
   2    |   group1   |   0
   3    |   group1   |   0

I need to update the state of a row by ID, while also returning some things:

I've got a query to do this as follows:

UPDATE t AS updated SET state = 1
FROM t as original
WHERE 
    updated.id = original.id AND
    updated.id = :some_id
RETURNING
    updated.state AS new_state,
    original.state AS old_state,
    (
        SELECT COUNT(*) FROM t 
        WHERE 
            group_name = updated.group_name AND
            state = 0
    ) as remaining_count;

However, it seems like the subquery within RETURNING is executed before the update has completed, leaving me with a remaining_count that is off by 1.

Additionally, I'm not sure how this behaves when concurrent queries are run. If we update two of these rows at the same time, is it possible that they would both return the same remaining_count?

Is there a more elegant solution to this? Perhaps some sort of window/aggregate function?

Upvotes: 5

Views: 9932

Answers (2)

Michael Ekoka
Michael Ekoka

Reputation: 20088

You could try (non-recursive) WITH queries, aka Common Table Expressions (CTEs). Their general structure is as follows:

WITH auxiliary_query_name AS (
    auxiliary_query_expression;
)
[, WITH ...] 
primary_query_expression;

Normally, auxiliary_query_expression and primary_query_expression run concurrently, and if they refer to the same underlying tables, the result is unpredictable. However, you can refer to auxiliary_query_name from within primary_query_expression, and from other auxiliary queries, thus enforcing a run sequence, where the referring query has to wait for the referred one to complete. Some finer points may apply, but that's the gist of it. CTEs also come with the advantage of being computed only once.

Regarding your query specifically, assuming that what you want in the end is the ID of the updated item, its old state, new state, the group it belongs to, and how many other items of that group are left to update, I believe the following would achieve this. I slightly modified the original query to update multiple items at once, to show how this approach shines (beside the clear sequence, it's performance advantages are moot if you update only a single item at a time).

WITH updated_t AS (
    UPDATE t AS updated SET state = 1
    FROM t as original
    WHERE 
        updated.id = original.id AND
        updated.id in :array_of_IDs -- I changed this 
    RETURNING
        updated.id,
        original.state AS old_state,
        updated.state AS new_state,
        updated.group_name
),
WITH remaining AS (
    SELECT t.group_name, count(*) as remaining_count
    -- we need to JOIN then filter out the updated rows because
    -- all WITH in a statement share the same snapshot, thus have
    -- the same starting "view" of base tables.
    FROM t LEFT JOIN updated_t
        ON t.id = updated_t.id
    WHERE updated_t.id is NULL
        AND t.group_name in (SELECT DISTINCT group_name from updated_t)
        AND t.state = 0
    GROUP BY group_name
)
SELECT 
    updated_t.id, 
    updated_t.group_name, 
    updated_t.old_state, 
    updated_t.new_state, 
    remaining.remaining_count
FROM updated_t, remaining
WHERE 
    updated_t.group_name = remaining.group_name;

Upvotes: 1

AdamKG
AdamKG

Reputation: 14066

The subquery is indeed run without seeing the change from the UPDATE, because it's running before the UPDATE has committed, and therefore it's not visible. Nevertheless, it's an easy fix; just add a where clause to filter out the ID you just updated in the subquery, making your query something like this:

UPDATE t AS updated SET state = 1
FROM t as original
WHERE 
    updated.id = original.id AND
    updated.id = :some_id
RETURNING
    updated.state AS new_state,
    original.state AS old_state,
    (
        SELECT COUNT(*) FROM t 
        WHERE 
            group_name = updated.group_name AND
            state = 0 AND
            t.id <> :some_id /* this is what I changed */
    ) as remaining_count;

Concurrency-wise, I'm not sure what the behavior would be, TBH; best I can do is point you at the relevant docs.

Upvotes: 7

Related Questions