update a table from insert result

Question

For the table B, there is a column named a_id which is the id of the table A. So a_id is the foreign key pointing to table a, but it is just a integer column and has no foreign constraint is set on it.

For each row in table B, we need to give the column a_id an integer value by creating a new record in table A.

The goal to do all below steps in one SQL.

1. Insert all of data into table A:

   insert into table A (name) values ('abc'), ('def'), ... returning id

2. Buck update a_id of each row in table B with the id(each id should be used once only) returned from step 1

   update table B set a_id = id(from previous insert statement)

Have tried something like:

update table B set a_id = (select ia.id from ia
(insert into table A (name) values ('abc'), ('def'), ... returning id) as ia)

But this gives a syntax error ERROR: syntax error at or near "into".

How to do it with one SQL?

Answer 1

If you only insert a single row, you can use a data modifying CTE

with new_row as (
  insert into table_A (name) values ('abc')
  returning id
)
update table_b
  set a_id = (select id from new_row)
where ?????; -- surely you don't want to update all rows in table_b

However the above will fail if your insert more than one row in the first statement.

It is unclear to me, which rows in table_b would need to be updated in that case.

Answer 2

This is not exactly one query, but if your concern is really just to avoid race conditions/transaction isolation anomalies resulting from a naive approach of doing this with multiple queries, then this should do it:

-- assign tableA's future primary keys into a_id without 
-- creating entries on tableA, pg_get_serial_sequence()
-- should retrieve tableA's primary key generator
UPDATE tableB
    SET a_id = nextval(pg_get_serial_sequence('tableA', 'id'))
    WHERE a_id IS NULL;

-- insert into tableB, with explicitly specified primary keys, 
-- note that this doesn't increment tableA's sequence as that 
-- has already been done in the previous operation by nextval()
INSERT INTO tableA(id, name)
    SELECT a_id, name FROM tableB;

Note that you should still wrap this in a transaction to ensure that tableB's updates gets rolled back if the insertion fails for any reason. Alternatively, since the two operations above are idempotent, they are safe to retry concurrently even without a transaction.