Bash and SQL scripts for testing plproxy server configuration

create_functions.sql

drop function if exists make_new_user(TEXT);

create function make_new_user(p_username TEXT, user_id OUT INTEGER ) as $$
begin
    INSERT INTO users (username, shard, last_updated) VALUES (p_username, (hashtext(p_username)&15), now())
    RETURNING users.user_id INTO user_id;
end;
$$ language plpgsql;

drop function if exists get_all_users_with_db();

create function get_all_users_with_db() returns setof users as $$
BEGIN
    return query
    select u.user_id, cast(current_database() as text), u.shard, u.last_updated from (
    SELECT * FROM users) as u;
end;
$$ language plpgsql;


drop function if exists get_all_users();

create function get_all_users() returns setof users as $$
BEGIN
    return query
    SELECT * FROM users;
end;
$$ language plpgsql;


drop function if exists get_user(TEXT);

create function get_user(p_username TEXT) returns setof users as $$
BEGIN
    return query
    select u.user_id, cast(current_database() as text), u.shard, u.last_updated from (
    SELECT * FROM users where username = p_username) as u;
end;
$$ language plpgsql;

create_proxy_functions.sql

drop function if exists make_new_user(TEXT);

create function make_new_user(p_username TEXT, user_id OUT INTEGER) as $$
    CLUSTER 'appplproxy';
    RUN ON hashtext(p_username);
$$ language plproxy;

drop function if exists get_all_users();

create function get_all_users() returns setof users as $$
    CLUSTER 'appplproxy';
    RUN ON ALL;
$$ language plproxy;

drop function if exists get_all_users_with_db();

create function get_all_users_with_db() returns setof users as $$
    CLUSTER 'appplproxy';
    RUN ON ALL;
$$ language plproxy;

drop function if exists get_user(TEXT);

create function get_user(p_username TEXT) returns setof users as $$
    CLUSTER 'appplproxy';
    RUN ON hashtext(p_username);
$$ language plproxy;

create_server_no_padding.sql

DROP SERVER appplproxy CASCADE;
CREATE SERVER appplproxy FOREIGN DATA WRAPPER plproxy
OPTIONS (connection_lifetime '1800',
         p0 'dbname=shard_p_0 host=127.0.0.1 port=5432',
         p1 'dbname=shard_p_1 host=127.0.0.1 port=5432',
         p2 'dbname=shard_p_2 host=127.0.0.1 port=5432',
         p3 'dbname=shard_p_3 host=127.0.0.1 port=5432',
         p4 'dbname=shard_p_4 host=127.0.0.1 port=5432',
         p5 'dbname=shard_p_5 host=127.0.0.1 port=5432',
         p6 'dbname=shard_p_6 host=127.0.0.1 port=5432',
         p7 'dbname=shard_p_7 host=127.0.0.1 port=5432',
         p8 'dbname=shard_p_8 host=127.0.0.1 port=5432',
         p9 'dbname=shard_p_9 host=127.0.0.1 port=5432',
         p10 'dbname=shard_p_10 host=127.0.0.1 port=5432',
         p11 'dbname=shard_p_11 host=127.0.0.1 port=5432',
         p12 'dbname=shard_p_12 host=127.0.0.1 port=5432',
         p13 'dbname=shard_p_13 host=127.0.0.1 port=5432',
         p14 'dbname=shard_p_14 host=127.0.0.1 port=5432',
         p15 'dbname=shard_p_15 host=127.0.0.1 port=5432'
);
CREATE USER MAPPING FOR plproxy_test SERVER appplproxy
        OPTIONS (user 'plproxy_test', password '123');
GRANT USAGE ON FOREIGN SERVER appplproxy TO plproxy_test;

create_server_with_padding.sql

DROP SERVER appplproxy CASCADE;
CREATE SERVER appplproxy FOREIGN DATA WRAPPER plproxy
OPTIONS (connection_lifetime '1800',
         p000 'dbname=shard_p_0 host=127.0.0.1 port=5432',
         p001 'dbname=shard_p_1 host=127.0.0.1 port=5432',
         p002 'dbname=shard_p_2 host=127.0.0.1 port=5432',
         p003 'dbname=shard_p_3 host=127.0.0.1 port=5432',
         p004 'dbname=shard_p_4 host=127.0.0.1 port=5432',
         p005 'dbname=shard_p_5 host=127.0.0.1 port=5432',
         p006 'dbname=shard_p_6 host=127.0.0.1 port=5432',
         p007 'dbname=shard_p_7 host=127.0.0.1 port=5432',
         p008 'dbname=shard_p_8 host=127.0.0.1 port=5432',
         p009 'dbname=shard_p_9 host=127.0.0.1 port=5432',
         p010 'dbname=shard_p_10 host=127.0.0.1 port=5432',
         p011 'dbname=shard_p_11 host=127.0.0.1 port=5432',
         p012 'dbname=shard_p_12 host=127.0.0.1 port=5432',
         p013 'dbname=shard_p_13 host=127.0.0.1 port=5432',
         p014 'dbname=shard_p_14 host=127.0.0.1 port=5432',
         p015 'dbname=shard_p_15 host=127.0.0.1 port=5432'
);
CREATE USER MAPPING FOR plproxy_test SERVER appplproxy
        OPTIONS (user 'plproxy_test', password '123');
GRANT USAGE ON FOREIGN SERVER appplproxy TO plproxy_test;

create_users_table.sql

drop table users;
create table users (
    user_id serial primary key,
    username text unique,
    shard text,
    last_updated timestamptz
);

plproxy test

# create user
psql -U commcarehq -h localhost -c "create user plproxy_test with password '123'"
psql -U commcarehq -h localhost -c "ALTER USER plproxy_test WITH SUPERUSER"

# create proxy database
psql -U plproxy_test -h localhost -d postgres -c "create database plproxy_test with owner plproxy_test"
psql -U plproxy_test -h localhost -d plproxy_test -c "CREATE EXTENSION plproxy"
psql -U plproxy_test -h localhost -d plproxy_test -f create_server_no_padding.sql
psql -U plproxy_test -h localhost -d plproxy_test -f create_users_table.sql
psql -U plproxy_test -h localhost -d plproxy_test -f create_proxy_functions.sql

# create shard databases
for a in p_{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}; do psql -U plproxy_test -h localhost -d postgres -c "create database shard_$a with owner plproxy_test"; done
for a in p_{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}; do psql -U plproxy_test -h localhost -d shard_$a -f create_users_table.sql; done
for a in p_{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}; do psql -U plproxy_test -h localhost -d shard_$a -f create_functions.sql; done

# populate each shards
for a in {1,2,3,4,5,8,10,11,13,14,16,17,23}; do psql -U plproxy_test -h localhost -c "select * from make_new_user('$a')"; done

# check the data
psql -U plproxy_test -h localhost -c "select * from get_all_users()"
psql -U plproxy_test -h localhost -c "select * from get_all_users_with_db()"
for a in {1,2,3,4,5,8,10,11,13,14,16,17,23}; do psql -U plproxy_test -h localhost -c "select * from get_user('$a')"; done

# now change the server config
psql -U plproxy_test -h localhost -d plproxy_test -f create_server_with_padding.sql

# make sure we can still fetch all the users by their username (hash key)
psql -U plproxy_test -h localhost -c "select * from get_all_users()"
psql -U plproxy_test -h localhost -c "select * from get_all_users_with_db()"
for a in {1,2,3,4,5,8,10,11,13,14,16,17,23}; do psql -U plproxy_test -h localhost -c "select * from get_user('$a')"; done

# cleanup
psql -U commcarehq -h localhost -d postgres -c "drop database plproxy_test"
for a in p_{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}; do psql -U plproxy_test -h localhost -d postgres -c "drop database shard_$a"; done
psql -U commcarehq -h localhost -d postgres -c "drop user plproxy_test"