Skip to content

Instantly share code, notes, and snippets.

@MaxGabriel
Created June 2, 2023 19:56
Show Gist options
  • Star 3 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save MaxGabriel/f8b4fcc7ed773dda79c62d69742d40d1 to your computer and use it in GitHub Desktop.
Save MaxGabriel/f8b4fcc7ed773dda79c62d69742d40d1 to your computer and use it in GitHub Desktop.
Postgres database wiping in tests, only deleting from tables with rows, using triggers to track inserts
CREATE TABLE tables_with_insertions (
name TEXT PRIMARY KEY
);
CREATE OR REPLACE FUNCTION track_insert_for_tests()
RETURNS TRIGGER AS $$
BEGIN
INSERT INTO tables_with_insertions VALUES (TG_TABLE_NAME) ON CONFLICT ON CONSTRAINT tables_with_insertions_pkey DO NOTHING;
RETURN NEW;
END;
$$ LANGUAGE 'plpgsql'
;
CREATE OR REPLACE FUNCTION add_insertion_tracking_triggers()
RETURNS void AS
$$
DECLARE
_tbl text;
_sch text;
BEGIN
FOR _sch, _tbl IN
SELECT schemaname, tablename
FROM pg_tables
WHERE schemaname = 'public'
AND NOT EXISTS (
SELECT true
FROM "schema_partial_wipe"
WHERE "table_name" = tablename
)
AND tablename != 'schema_partial_wipe'
AND tablename != 'tables_with_insertions'
LOOP
EXECUTE format('CREATE TRIGGER %I_track_inserts_for_tests AFTER INSERT ON %I.%I FOR EACH STATEMENT EXECUTE PROCEDURE track_insert_for_tests()', _tbl, _sch, _tbl);
END LOOP;
END
$$ LANGUAGE 'plpgsql'
;
SELECT add_insertion_tracking_triggers();
-- Delete data from all tables besides those whitelisted in schema_partial_wipe
-- If you're inserting data in a migration and expecting it to be present in a test,
-- you should either insert the data in your test or add it to schema_partial_wipe
CREATE OR REPLACE FUNCTION f_delete_tables()
RETURNS void AS
$func$
DECLARE
_tbl text;
_sch text;
BEGIN
-- disable all triggers
SET session_replication_role = replica;
-- defer all constraints to the end
SET CONSTRAINTS ALL DEFERRED;
-- delete all rows from all tables that had data inserted by the most recent tests
FOR _tbl IN
SELECT name
FROM tables_with_insertions
WHERE name != 'schema_partial_wipe'
AND NOT EXISTS (
SELECT true
FROM "schema_partial_wipe"
WHERE "table_name" = tables_with_insertions.name
)
LOOP
EXECUTE format('DELETE FROM public.%I', _tbl);
END LOOP;
FOR _sch, _tbl IN
SELECT schemaname, tablename
FROM pg_tables
JOIN schema_partial_wipe spw
ON spw.table_name = tablename
WHERE schemaname = 'public'
AND tablename != 'schema_migrations'
LOOP
EXECUTE format('DELETE FROM %I.%I WHERE "dont_wipe" = false', _sch, _tbl);
END LOOP;
DELETE FROM tables_with_insertions;
-- re-enable all triggers now that we're done
SET session_replication_role = DEFAULT;
END
$func$ LANGUAGE plpgsql;
-- f_delete_tables only deletes tables with data and runs in between tests
-- This function lets us wipe all data (besides that in schema_partial wipe)
-- We do this once at the start of the test suite. This means that data that's
-- not whitelisted by schema_partial_wipe will always be deleted at the start
-- of a test, regardless of test execution order.
CREATE OR REPLACE FUNCTION f_delete_all_tables()
RETURNS void AS
$$
DECLARE
_tbl text;
_sch text;
BEGIN
FOR _sch, _tbl IN
SELECT schemaname, tablename
FROM pg_tables
WHERE schemaname = 'public'
AND NOT EXISTS (
SELECT true
FROM "schema_partial_wipe"
WHERE "table_name" = tablename
)
AND tablename != 'schema_partial_wipe'
AND tablename != 'tables_with_insertions'
AND tablename != 'schema_migrations'
LOOP
EXECUTE format('DELETE FROM %I.%I', _sch, _tbl);
END LOOP;
END
$$ LANGUAGE 'plpgsql'
;
SELECT f_delete_all_tables();
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment