PostgreSQL Atomic Swaps in a Function

postgres-atomic-swap.sql

-- --------------------------------------------------------
-- -- Table: UHS
-- --   Creates a table called 'uhs' that holds UUIDs.
-- --------------------------------------------------------

CREATE TABLE uhs (
  id UUID NOT NULL UNIQUE DEFAULT gen_random_uuid(),
  CONSTRAINT pk_uhs_id PRIMARY KEY (id)
) WITH (OIDS=FALSE);

CREATE INDEX idx_uhs_id ON uhs USING btree (id);

-- --------------------------------------------------------
-- -- Function: swap
-- --   Creates a function that DELETEs a set of UUIDs and
-- --   replaces them with a new set of UUIDs atomically.
-- --   If there are any errors in DELETE or INSERT
-- --   everything rolls back as if nothing had happened.
-- --------------------------------------------------------

CREATE OR REPLACE FUNCTION swap(_delete_ids UUID[], _create_ids UUID[])
RETURNS VOID AS $$
DECLARE
  count INTEGER := 0;
BEGIN
  DELETE FROM uhs WHERE id = ANY (_delete_ids);
  GET DIAGNOSTICS count = ROW_COUNT;
  IF count = cardinality(_delete_ids) THEN -- if all the deletes work, try the inserts
    INSERT INTO uhs (id) VALUES  (UNNEST (_create_ids));
  ELSE
    RAISE NOTICE 'Not all UHS IDs exist. (missing %)', cardinality(_delete_ids) - count;
  END IF;
END;
$$
LANGUAGE plpgsql;

-- --------------------------------------------------------
-- Add some test data
-- --------------------------------------------------------

INSERT INTO uhs
  (id)
VALUES
  ('d335ab4f-063b-4a7d-88d4-4f8e835db7b0'::uuid),
  ('be33395e-3c25-447c-a3f5-6dfee29f18bd'::uuid);

-- --------------------------------------------------------
-- Do a swap
-- --------------------------------------------------------

SELECT swap(
  ARRAY[
    'd335ab4f-063b-4a7d-88d4-4f8e835db7b0'::uuid,
    'be33395e-3c25-447c-a3f5-6dfee29f18bd'::uuid
  ],
  ARRAY[
    '7f58fde0-e742-474b-963e-53a42dc1049f'::uuid,
    '6cb24995-dc87-4fa6-b1d3-47e1ede95d97'::uuid
  ]
);