testing_drop_index_on_partitioned_table.sql

-- Reset the database (! BE CAREFUL, DON'T DO THIS ON A WORKING DB !)
DROP SCHEMA public CASCADE;
CREATE SCHEMA public;

-- Helper for testing the acquired locks
-- Usage:
--   :locks
\set locks '
  SELECT pid,
         locktype,
         CASE locktype
            WHEN ''relation'' THEN relation::regclass::text
            WHEN ''transactionid'' THEN transactionid::text
            WHEN ''virtualxid'' THEN virtualxid
         END AS lockid,
         mode,
         granted
  FROM pg_locks
  ORDER BY 5, 4;'

CREATE TABLE public.cities_partitioned (
  id bigint NOT NULL,
  name varchar,
  country_id bigint
) PARTITION BY RANGE (country_id);

CREATE SEQUENCE public.cities_id_seq START WITH 1 INCREMENT BY 1 NO MINVALUE NO MAXVALUE CACHE 1;
ALTER SEQUENCE public.cities_id_seq OWNED BY public.cities_partitioned.id;
ALTER TABLE public.cities_partitioned ALTER COLUMN id SET DEFAULT nextval('public.cities_id_seq'::regclass);

ALTER TABLE public.cities_partitioned
    ADD CONSTRAINT cities_partitioned_pkey PRIMARY KEY (id, country_id);
CREATE INDEX index_cities_p_on_name ON public.cities_partitioned USING btree (name);

CREATE TABLE public.countries (
  id bigint NOT NULL,
  name varchar
);

CREATE SEQUENCE public.countries_id_seq START WITH 1 INCREMENT BY 1 NO MINVALUE NO MAXVALUE CACHE 1;
ALTER SEQUENCE public.countries_id_seq OWNED BY public.countries.id;
ALTER TABLE public.countries ALTER COLUMN id SET DEFAULT nextval('public.countries_id_seq'::regclass);
ALTER TABLE public.countries
    ADD CONSTRAINT countries_pkey PRIMARY KEY (id);

ALTER TABLE public.cities_partitioned
    ADD CONSTRAINT fk_cities_to_county FOREIGN KEY (country_id) REFERENCES public.countries(id) ON DELETE CASCADE;

CREATE OR REPLACE FUNCTION public.create_cities_partitions() RETURNS trigger LANGUAGE plpgsql AS $$
  DECLARE
    current_id bigint;
    partition_index bigint;
    partition_size bigint;
    lower bigint;
    upper bigint;
  BEGIN
    current_id := COALESCE(NEW.id, CURRVAL('public.countries_id_seq'::regclass));
    IF NOT (EXISTS (SELECT * FROM pg_tables WHERE schemaname = 'public' AND tablename = 'cities_' || current_id)) THEN
      EXECUTE 'CREATE TABLE IF NOT EXISTS public.cities_' || current_id || ' (
        LIKE public.cities_partitioned
          INCLUDING DEFAULTS
          INCLUDING CONSTRAINTS
          INCLUDING GENERATED
          INCLUDING INDEXES
          INCLUDING STATISTICS
        );';
      EXECUTE 'ALTER TABLE public.cities_' || current_id || ' ADD CONSTRAINT cities_partition_' || current_id || ' CHECK (country_id = ' || current_id || ');';
      EXECUTE 'ALTER TABLE public.cities_partitioned ATTACH PARTITION public.cities_' || current_id || ' FOR VALUES FROM (' || current_id || ') TO (' || current_id + 1 || ');';
      EXECUTE 'ALTER TABLE public.cities_' || current_id || ' DROP CONSTRAINT cities_partition_' || current_id || ';';
    END IF;
    RETURN NEW;
  END;
$$;

CREATE TRIGGER auto_partition BEFORE INSERT OR UPDATE ON public.countries FOR EACH ROW EXECUTE FUNCTION public.create_cities_partitions();

CREATE TABLE public.shops (
  id bigint NOT NULL,
  name varchar,
  country_id bigint NOT NULL,
  city_id bigint NOT NULL
);

CREATE SEQUENCE public.shops_id_seq START WITH 1 INCREMENT BY 1 NO MINVALUE NO MAXVALUE CACHE 1;
ALTER SEQUENCE public.shops_id_seq OWNED BY public.shops.id;
ALTER TABLE public.shops ALTER COLUMN id SET DEFAULT nextval('public.shops_id_seq'::regclass);
ALTER TABLE public.shops
    ADD CONSTRAINT fk_shops_to_city_and_county FOREIGN KEY (country_id, city_id) REFERENCES public.cities_partitioned(country_id, id) ON DELETE CASCADE;

INSERT INTO countries (name) VALUES ('US'), ('UK'), ('CA'), ('AU');
INSERT INTO cities_partitioned (name, country_id) VALUES
  ('New York', 1),
  ('Boston', 1),
  ('Los Angeles', 1),
  ('Dallas', 1),
  ('Austin', 1),
  ('Seattle', 1),
  ('Washington', 1),
  ('London', 2),
  ('Brighton', 2),
  ('Leeds', 2),
  ('Bristol', 2),
  ('Vancouver', 3),
  ('Alberta', 3),
  ('Ontario', 3),
  ('Yukon', 3),
  ('Sydney', 4),
  ('Melbourne', 4),
  ('Darwin', 4),
  ('Brisbane', 4);

INSERT INTO shops (name, city_id, country_id) VALUES
  ('H&M', 1, 1),
  ('Ikea', 8, 2),
  ('Gallery', 12, 3),
  ('Ostin', 16, 4);

--- Dropping the index

-- In a separate connection start a transaction wich AccessShareLock on cities_partitioned
-- BEGIN;
--   SELECT * from cities_partitioned;

-- Take 1: trying to detach all partitions, drop the index, then attach partitions again
-- FAILED: dropping a constraint always acquires AccessExclusiveLock
ALTER TABLE shops DROP CONSTRAINT fk_shops_to_city_and_county;
ALTER TABLE cities_partitioned DETACH PARTITION cities_4 CONCURRENTLY;
ALTER TABLE cities_partitioned DETACH PARTITION cities_3 CONCURRENTLY;
ALTER TABLE cities_partitioned DETACH PARTITION cities_2 CONCURRENTLY;
ALTER TABLE cities_partitioned DETACH PARTITION cities_1 CONCURRENTLY;
DROP INDEX index_cities_p_on_name;
ALTER TABLE cities_partitioned ATTACH PARTITION cities_1 FOR VALUES FROM (1) TO (2);
ALTER TABLE cities_partitioned ATTACH PARTITION cities_2 FOR VALUES FROM (2) TO (3);
ALTER TABLE cities_partitioned ATTACH PARTITION cities_3 FOR VALUES FROM (3) TO (4);
ALTER TABLE cities_partitioned ATTACH PARTITION cities_4 FOR VALUES FROM (4) TO (5);
ALTER TABLE public.shops
  ADD CONSTRAINT fk_shops_to_city_and_county FOREIGN KEY (country_id, city_id) REFERENCES public.cities_partitioned(country_id, id) ON DELETE CASCADE NOT VALID;
ALTER TABLE public.shops VALIDATE CONSTRAINT fk_shops_to_city_and_county;

DROP INDEX CONCURRENTLY cities_1_name_idx;
DROP INDEX CONCURRENTLY cities_2_name_idx;
DROP INDEX CONCURRENTLY cities_3_name_idx;
DROP INDEX CONCURRENTLY cities_4_name_idx;

--- Create the index again
CREATE INDEX index_cities_p_on_name ON public.cities_partitioned USING btree (name);