citus-pgbench-distribute.sql

Original query:

/* ==== STEP 1. Drop FKs ===== */

-- accounts -> branches.
ALTER TABLE pgbench_accounts DROP CONSTRAINT pgbench_accounts_bid_fkey;
-- tellers -> branches.
ALTER TABLE pgbench_tellers DROP CONSTRAINT pgbench_tellers_bid_fkey;
-- history -> branches.
ALTER TABLE pgbench_history DROP CONSTRAINT pgbench_history_bid_fkey;
-- history -> accounts.
ALTER TABLE pgbench_history DROP CONSTRAINT pgbench_history_aid_fkey;
-- history -> tellers.
ALTER TABLE pgbench_history DROP CONSTRAINT pgbench_history_tid_fkey;

/* ==== STEP 2. modifications to schema ===== */
ALTER TABLE pgbench_history ADD CONSTRAINT pgbench_history_pkey PRIMARY KEY (bid, tid, aid, mtime);

/* NOTE: to distrubute by "bid", more modifications are needed. PKeys should include "bid" column.
ALTER TABLE pgbench_accounts DROP CONSTRAINT pgbench_accounts_pkey CASCADE;
ALTER TABLE pgbench_accounts ADD CONSTRAINT pgbench_accounts_pkey PRIMARY KEY (bid, aid);
-- This is forbidden because of Citus limitations: CREATE UNIQUE INDEX ON pgbench_accounts (aid);
CREATE INDEX ON pgbench_accounts (aid);

ALTER TABLE pgbench_tellers DROP CONSTRAINT pgbench_tellers_pkey CASCADE;
ALTER TABLE pgbench_tellers ADD CONSTRAINT pgbench_tellers_pkey PRIMARY KEY (bid, tid);
-- This is forbidden because of Citus limitations: CREATE UNIQUE INDEX ON pgbench_tellers (tid);
CREATE INDEX ON pgbench_tellers (tid); */

/* ===== STEP 3. Distribute tables ===== */
-- SET citus.shard_count = 32;
SELECT create_reference_table('pgbench_branches');
SELECT create_reference_table('pgbench_tellers');
SELECT create_distributed_table('pgbench_accounts', 'aid');
SELECT create_distributed_table('pgbench_history', 'aid');

/* ===== STEP 4. Create FKs back */
-- accounts -> branches.
ALTER TABLE pgbench_accounts ADD CONSTRAINT pgbench_accounts_bid_fkey FOREIGN KEY (bid) REFERENCES pgbench_branches(bid);

-- tellers -> branches.
ALTER TABLE pgbench_tellers ADD CONSTRAINT pgbench_tellers_bid_fkey FOREIGN KEY (bid) REFERENCES pgbench_branches(bid);

-- history -> branches.
ALTER TABLE pgbench_history ADD CONSTRAINT pgbench_history_bid_fkey FOREIGN KEY (bid) REFERENCES pgbench_branches(bid);

-- history -> accounts. When distributing by bid, this is forbidden because of Citus limitations.
ALTER TABLE pgbench_history ADD CONSTRAINT pgbench_history_aid_fkey FOREIGN KEY (aid) REFERENCES pgbench_accounts(aid);

-- history -> tellers. When distributing by bid, this is forbidden because of Citus limitations.
ALTER TABLE pgbench_history ADD CONSTRAINT pgbench_history_tid_fkey FOREIGN KEY (tid) REFERENCES pgbench_tellers(tid);


Prettified:

ALTER TABLE "pgbench_accounts" DROP CONSTRAINT pgbench_accounts_bid_fkey; ALTER TABLE "pgbench_tellers" DROP CONSTRAINT pgbench_tellers_bid_fkey; ALTER TABLE "pgbench_history" DROP CONSTRAINT pgbench_history_bid_fkey; ALTER TABLE "pgbench_history" DROP CONSTRAINT pgbench_history_aid_fkey; ALTER TABLE "pgbench_history" DROP CONSTRAINT pgbench_history_tid_fkey; ALTER TABLE "pgbench_history" ADD CONSTRAINT pgbench_history_pkey PRIMARY KEY ("bid", "tid", "aid", "mtime"); SELECT
  create_reference_table('pgbench_branches'); SELECT
  create_reference_table('pgbench_tellers'); SELECT
  create_distributed_table('pgbench_accounts', 'aid'); SELECT
  create_distributed_table('pgbench_history', 'aid'); ALTER TABLE "pgbench_accounts" ADD CONSTRAINT pgbench_accounts_bid_fkey FOREIGN KEY ("bid") REFERENCES "pgbench_branches" ("bid"); ALTER TABLE "pgbench_tellers" ADD CONSTRAINT pgbench_tellers_bid_fkey FOREIGN KEY ("bid") REFERENCES "pgbench_branches" ("bid"); ALTER TABLE "pgbench_history" ADD CONSTRAINT pgbench_history_bid_fkey FOREIGN KEY ("bid") REFERENCES "pgbench_branches" ("bid"); ALTER TABLE "pgbench_history" ADD CONSTRAINT pgbench_history_aid_fkey FOREIGN KEY ("aid") REFERENCES "pgbench_accounts" ("aid"); ALTER TABLE "pgbench_history" ADD CONSTRAINT pgbench_history_tid_fkey FOREIGN KEY ("tid") REFERENCES "pgbench_tellers" ("tid")