Postgres Multiple WITH blocks to do multiple inserts and updates in one query

Question

My database has 3 tables, users, firms, and stashes. The stashes belong to users and the users belong to firms. The users also have a foreign key for currently active stash.

I want to run a query when a new user registers to create a new firm for them, create a new user, create a new stash, and also set that stash to be active. Here is what I have and it runs but it is not setting the active stash parameter on the new user correctly.

I would love some help and also just general advice if the way that I have structured this query is not the most efficient. Thank you

WITH insert1 AS (
    INSERT INTO firms (name) VALUES ('Jack_testFirm') RETURNING id AS ret_firm_id
)
,
insert2 AS (
INSERT INTO users (email, admin, firm_id) 
SELECT 'jack@gmail.com', TRUE, ret_firm_id 
    FROM insert1
RETURNING users.id AS new_user_id
)
,
insert3 AS (
INSERT INTO stashes (name, user_id)
SELECT 'Stash 1', new_user_id FROM insert2
RETURNING stashes.id AS new_stash_id
)

UPDATE users
SET active_stash = (SELECT new_stash_id FROM insert3)
WHERE users.id = (SELECT new_user_id FROM insert2)

After the query, I should have a brand new user, 'jack@gmail.com' which belongs to the firm 'Jack_testFirm', a new stash which belongs to 'jack@gmail.com' and jack@gmail.com should store that stashes id as the active stash in the users table.

You cannot modify the same row twice in a single query. In this case, both `insert2` and the main `update` are working on the same row. — Mike Organek, Aug 26 '22 at 01:18
@MikeOrganek how can I store the results of the first query to use in the UPDATE row query? I want to store the ids of the newly created rows — Ari Baranian, Aug 26 '22 at 01:28
I don't think you can do that In one step. Because, all the inserts need to be materialized (insert to the table) then do the update. You can wrap it in an transaction. That means all fail or all success. — jian, Aug 26 '22 at 07:57

jian · Accepted Answer · 2022-08-26T09:50:01.290

Use procedure. Idea is first materialize all of your insert operation, use variable hold your new_stash_id and new_user_id, then do the update.
Procedure all happen in a single transaction (Do stored procedures run in database transaction in Postgres?), all fail or all success.
dbfiddle

Call procedure: call test()

Code:

CREATE OR REPLACE PROCEDURE test ()
LANGUAGE plpgsql
AS $$
DECLARE
    _new_stash_id bigint;
    _new_user_id bigint;
BEGIN
    WITH insert1 AS (
INSERT INTO firms (name)
            VALUES ('Jack_testFirm')
        RETURNING
            firm_id AS ret_firm_id
), insert2 AS (
INSERT INTO users (email, admin, firm_id)
    SELECT
        'jack@gmail.com',
        TRUE,
        ret_firm_id
    FROM
        insert1
    RETURNING
        users.user_id AS new_user_id
),
insert3 AS (
INSERT INTO stashes (name, user_id)
    SELECT
        'Stash 1',
        new_user_id
    FROM
        insert2
    RETURNING
        new_stash_id,
        user_id
)
SELECT
    new_stash_id,
    user_id
FROM
    insert3 INTO _new_stash_id,
    _new_user_id;
        RAISE NOTICE '_new_stash_id: %', _new_stash_id;
        RAISE NOTICE '_new_user_id: %', _new_user_id;
        UPDATE
            users
        SET
            active_stash = _new_stash_id
        WHERE
            user_id = _new_user_id;
END
$$;

Postgres Multiple WITH blocks to do multiple inserts and updates in one query

1 Answers1