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.

Answer 1

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
$$;