Local Citus cluster setup

local-citus.sh

#!/bin/bash

role=citus
database=citus
port=5432
workers=2
timeout=30

# start cluster with 2 workers
[ -f docker-compose.yml ] || wget https://raw.githubusercontent.com/citusdata/docker/master/docker-compose.yml
docker-compose -p citus up --scale worker=${workers} -d

# ensure the cluster is fully running
t=1
while ! grep "^${workers}$"<<<"$( psql postgres://postgres@localhost:${port}?sslmode=disable -c "COPY(select count(*) from master_get_active_worker_nodes()) TO STDOUT DELIMITER ','" )";
do
    timeout=$(( ${timeout} + 1 ))
    if [ "${t}" == "${timeout}" ]
    then
        echo "Timeout: could not find ${workers} workers attached. Exiting..."
        exit 1
    fi
    sleep 1
done

# create role and database
psql postgres://postgres@localhost:${port}?sslmode=disable <<EOQ
  CREATE ROLE ${role} WITH LOGIN IN ROLE pg_monitor;
  SELECT run_command_on_workers('CREATE ROLE ${role} WITH LOGIN IN ROLE pg_monitor');

  CREATE DATABASE ${database} WITH OWNER ${role};
  SELECT run_command_on_workers('CREATE DATABASE ${database} WITH OWNER=${role}');
EOQ

# install extensions for database
psql postgres://postgres@localhost:${port}/${database}?sslmode=disable <<EOQ
  CREATE EXTENSION IF NOT EXISTS "citus";
  CREATE EXTENSION IF NOT EXISTS "hll";
  CREATE EXTENSION IF NOT EXISTS "topn";
EOQ

# attach workers to database
for ((wid=1 ; wid<=${workers} ; wid+=1));
do
    worker_ip="$( docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' citus_worker_${wid} )"

    psql postgres://postgres@localhost:${port}/${database}?sslmode=disable <<EOQ
      SELECT * from master_add_node('${worker_ip}', ${port});
EOQ
done

# install extensions on workers
psql postgres://postgres@localhost:${port}/${database}?sslmode=disable <<EOQ
 SELECT run_command_on_workers('CREATE EXTENSION IF NOT EXISTS "citus"');
 SELECT run_command_on_workers('CREATE EXTENSION IF NOT EXISTS "hll"');
 SELECT run_command_on_workers('CREATE EXTENSION IF NOT EXISTS "topn"');
EOQ