Instantly share code, notes, and snippets.

Embed
What would you like to do?
Local Citus cluster setup
#!/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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment