canaries.sql

-- Users

CREATE TABLE users(
  id bigint NOT NULL,
  name text,
  PRIMARY KEY (id)
);

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

-- Resources

CREATE TABLE resources(
  id bigint NOT NULL,
  user_id bigint NOT NULL,
  title character varying,
  PRIMARY KEY (id)
);

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

-- Data

INSERT INTO users (name) SELECT 'User ' || g.id FROM generate_series(1,10000) as g(id);
INSERT INTO resources (user_id, title) SELECT floor(random() * 10000 + 1), 'Resource ' || g.id FROM generate_series(1,5000000) as g(id);

-- Foreign Key

ALTER TABLE ONLY public.resources ADD CONSTRAINT fk_resources_user_id FOREIGN KEY (user_id) REFERENCES public.users(id);

-- Partitioned resources

CREATE TABLE public.resources_partitioned (
  id bigint NOT NULL,
  user_id bigint NOT NULL,
  title character varying,
  PRIMARY KEY (id, user_id)
) PARTITION BY RANGE (user_id);


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

-- Auto-partitioning (example)

CREATE OR REPLACE FUNCTION create_resource_partitions() RETURNS TRIGGER AS $auto_partition$
  DECLARE
    current_id bigint;
    partition_number bigint;
    partition_size bigint;
    lower bigint;
    upper bigint;
  BEGIN
    current_id        := COALESCE(NEW.id, CURRVAL('public.users_id_seq'::regclass));
    partition_size    := 1000;
    partition_number  := (current_id - 1) / partition_size + 1;

    -- This is for preventing redundant partition creation attempts
    -- IF (current_id % partition_size > partition_size - 100) THEN
    --   RETURN NEW;
    -- END IF;

    lower := 1 + partition_size * partition_number;
    upper := 1 + partition_size * (partition_number + 1);

    IF NOT (EXISTS(SELECT * FROM pg_tables WHERE schemaname = 'public' AND tablename = 'resources_' || upper - 1)) THEN
      EXECUTE 'CREATE TABLE IF NOT EXISTS public.resources_' || upper - 1 || ' (
                LIKE public.resources_partitioned
                INCLUDING DEFAULTS
                INCLUDING CONSTRAINTS
                INCLUDING GENERATED
                INCLUDING INDEXES
              );';

      EXECUTE 'ALTER TABLE public.resources_' || upper - 1 || ' ADD CONSTRAINT resources_partition_' || upper - 1  || '
        CHECK (user_id >= ' || lower || ' AND user_id < ' || upper || ');';

      EXECUTE 'ALTER TABLE public.resources_partitioned ATTACH PARTITION public.resources_' || upper - 1 || '
        FOR VALUES FROM (' || lower || ') TO (' || upper || ');';

      EXECUTE 'ALTER TABLE public.resources_' || upper - 1 || ' DROP CONSTRAINT resources_partition_' || upper - 1  || ';';
    END IF;

    RETURN NEW;
  END;
$auto_partition$ LANGUAGE plpgsql;

CREATE TRIGGER auto_partition BEFORE INSERT OR UPDATE ON public.users FOR EACH ROW EXECUTE PROCEDURE create_resource_partitions();

-- Initial partitions

CREATE OR REPLACE FUNCTION pg_temp.generate_initial_resource_partitions() RETURNS void AS $BODY$
  DECLARE
    current_id bigint;
    partition_number bigint;
    partition_size bigint;
    lower bigint;
    upper bigint;
  BEGIN
    partition_size := 1000;

    SELECT GREATEST(NEXTVAL('public.users_id_seq'), 999) INTO current_id;

    FOR partition_number IN SELECT * FROM generate_series(0, 1 + current_id / partition_size) LOOP
      lower := 1 + partition_size * partition_number;
      upper := 1 + partition_size * (partition_number + 1);

      EXECUTE 'CREATE TABLE IF NOT EXISTS public.resources_' || upper - 1 || ' PARTITION OF
        public.resources_partitioned FOR VALUES FROM (' || lower || ') TO (' || upper || ');';
    END LOOP;

    RETURN;
  END;
$BODY$ LANGUAGE plpgsql;

SELECT * FROM pg_temp.generate_initial_resource_partitions();

  
-- Sync trigges (resources -> resources_partitioned)

CREATE OR REPLACE FUNCTION public.resources_insert_trigger()
RETURNS TRIGGER LANGUAGE plpgsql AS $$
BEGIN
  INSERT INTO public.resources_partitioned (
    id,
    user_id,
    title
  )
  SELECT
    new.id,
    new.user_id,
    new.title
  ;
  RETURN new;
END $$;

CREATE TRIGGER resources_insert_trigger
  AFTER INSERT ON public.resources
  FOR EACH ROW WHEN (pg_trigger_depth() = 0)
  EXECUTE PROCEDURE public.resources_insert_trigger();

CREATE OR REPLACE FUNCTION public.resources_update_trigger()
RETURNS TRIGGER LANGUAGE plpgsql AS $$
BEGIN
  UPDATE public.resources_partitioned SET
    title = new.title
  WHERE id = new.id AND user_id = new.user_id;
  RETURN new;
END $$;

CREATE TRIGGER resources_update_trigger
  AFTER UPDATE ON public.resources
  FOR EACH ROW WHEN (pg_trigger_depth() = 0)
  EXECUTE PROCEDURE public.resources_update_trigger();

CREATE OR REPLACE FUNCTION public.resources_delete_trigger()
RETURNS TRIGGER LANGUAGE plpgsql AS $$
BEGIN
  DELETE FROM public.resources_partitioned WHERE id = old.id AND user_id = old.user_id;
  RETURN new;
END $$;

CREATE TRIGGER resources_delete_trigger
  AFTER DELETE ON public.resources
  FOR EACH ROW WHEN (pg_trigger_depth() = 0)
  EXECUTE PROCEDURE public.resources_delete_trigger();

-- Copy data (plain)

INSERT INTO resources_partitioned(id, user_id, title) SELECT id, user_id, title FROM resources;

-- Reverse-sync triggers (resources_partitioned -> resources)

CREATE OR REPLACE FUNCTION public.resources_partitioned_insert_trigger()
RETURNS TRIGGER LANGUAGE plpgsql AS $$
BEGIN
  INSERT INTO public.resources (
    id,
    user_id,
    title
  )
  SELECT
    new.id,
    new.user_id,
    new.title
  ;
  RETURN new;
END $$;

CREATE TRIGGER resources_partitioned_insert_trigger
  AFTER INSERT ON public.resources_partitioned
  FOR EACH ROW WHEN (pg_trigger_depth() = 0)
  EXECUTE PROCEDURE public.resources_partitioned_insert_trigger();

CREATE OR REPLACE FUNCTION public.resources_partitioned_update_trigger()
RETURNS TRIGGER LANGUAGE plpgsql AS $$
BEGIN
  UPDATE public.resources SET
    title = new.title
  WHERE id = new.id AND user_id = new.user_id;
  RETURN new;
END $$;

CREATE TRIGGER resources_partitioned_update_trigger
  AFTER UPDATE ON public.resources_partitioned
  FOR EACH ROW WHEN (pg_trigger_depth() = 0)
  EXECUTE PROCEDURE public.resources_partitioned_update_trigger();

CREATE OR REPLACE FUNCTION public.resources_partitioned_delete_trigger()
RETURNS TRIGGER LANGUAGE plpgsql AS $$
BEGIN
  DELETE FROM public.resources WHERE id = old.id AND user_id = old.user_id;
  RETURN new;
END $$;

CREATE TRIGGER resources_partitioned_delete_trigger
  AFTER DELETE ON public.resources_partitioned
  FOR EACH ROW WHEN (pg_trigger_depth() = 0)
  EXECUTE PROCEDURE public.resources_partitioned_delete_trigger();

-- TADA!

UPDATE resources SET title='test' WHERE id=42;
SELECT title FROM resources_partitioned WHERE id=42; #=> test

UPDATE resources_partitioned SET title='Resource 42' WHERE id=42;
SELECT title FROM resources WHERE id=42; #=> Resource 42