citus/src/test/regress/sql/local_shard_copy.sql

338 lines
7.5 KiB
PL/PgSQL

CREATE SCHEMA local_shard_copy;
SET search_path TO local_shard_copy;
SET client_min_messages TO DEBUG;
SELECT * FROM master_add_node('localhost', :master_port, groupid := 0);
SET citus.shard_count TO 4;
SET citus.shard_replication_factor TO 1;
SET citus.replication_model TO 'streaming';
CREATE TABLE reference_table (key int PRIMARY KEY);
SELECT create_reference_table('reference_table');
CREATE TABLE distributed_table (key int PRIMARY KEY, age bigint CHECK (age >= 10));
SELECT create_distributed_table('distributed_table','key');
INSERT INTO distributed_table SELECT *,* FROM generate_series(20, 40);
INSERT INTO reference_table SELECT * FROM generate_series(1, 10);
CREATE TABLE local_table (key int PRIMARY KEY);
INSERT INTO local_table SELECT * from generate_series(1, 10);
-- partitioned table
CREATE TABLE collections_list (
key bigserial,
collection_id integer
) PARTITION BY LIST (collection_id );
SELECT create_distributed_table('collections_list', 'key');
CREATE TABLE collections_list_0
PARTITION OF collections_list (key, collection_id)
FOR VALUES IN ( 0 );
CREATE TABLE collections_list_1
PARTITION OF collections_list (key, collection_id)
FOR VALUES IN ( 1 );
-- connection worker and get ready for the tests
\c - - - :worker_1_port
SET search_path TO local_shard_copy;
SET citus.log_local_commands TO ON;
-- returns true of the distribution key filter
-- on the distributed tables (e.g., WHERE key = 1), we'll hit a shard
-- placement which is local to this not
CREATE OR REPLACE FUNCTION shard_of_distribution_column_is_local(dist_key int) RETURNS bool AS $$
DECLARE shard_is_local BOOLEAN := FALSE;
BEGIN
WITH local_shard_ids AS (SELECT get_shard_id_for_distribution_column('local_shard_copy.distributed_table', dist_key)),
all_local_shard_ids_on_node AS (SELECT shardid FROM pg_dist_placement WHERE groupid IN (SELECT groupid FROM pg_dist_local_group))
SELECT
true INTO shard_is_local
FROM
local_shard_ids
WHERE
get_shard_id_for_distribution_column IN (SELECT * FROM all_local_shard_ids_on_node);
IF shard_is_local IS NULL THEN
shard_is_local = FALSE;
END IF;
RETURN shard_is_local;
END;
$$ LANGUAGE plpgsql;
-- pick some example values that reside on the shards locally and remote
-- distribution key values of 1,6, 500 and 701 are LOCAL to shards,
-- we'll use these values in the tests
SELECT shard_of_distribution_column_is_local(1);
SELECT shard_of_distribution_column_is_local(6);
SELECT shard_of_distribution_column_is_local(500);
SELECT shard_of_distribution_column_is_local(701);
-- distribution key values of 11 and 12 are REMOTE to shards
SELECT shard_of_distribution_column_is_local(11);
SELECT shard_of_distribution_column_is_local(12);
BEGIN;
-- run select with local execution
SELECT count(*) FROM distributed_table WHERE key = 1;
SELECT count(*) FROM distributed_table;
-- the local placements should be executed locally
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
-- verify that the copy is successful.
SELECT count(*) FROM distributed_table;
ROLLBACK;
BEGIN;
-- run select with local execution
SELECT count(*) FROM distributed_table WHERE key = 1;
SELECT count(*) FROM distributed_table;
-- the local placements should be executed locally
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
-- verify the put ages.
SELECT * FROM distributed_table;
ROLLBACK;
BEGIN;
-- run select with local execution
SELECT count(*) FROM distributed_table WHERE key = 1;
SELECT count(*) FROM distributed_table;
-- the local placements should be executed locally
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
-- verify that the copy is successful.
SELECT count(*) FROM distributed_table;
ROLLBACK;
BEGIN;
-- run select with local execution
SELECT age FROM distributed_table WHERE key = 1;
SELECT count(*) FROM collections_list;
-- the local placements should be executed locally
COPY collections_list FROM STDIN WITH delimiter ',';
1, 0
2, 0
3, 0
4, 1
5, 1
\.
-- verify that the copy is successful.
SELECT count(*) FROM collections_list;
ROLLBACK;
BEGIN;
-- run select with local execution
SELECT age FROM distributed_table WHERE key = 1;
SELECT count(*) FROM distributed_table;
-- the local placements should be executed locally
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
-- verify that the copy is successful.
SELECT count(*) FROM distributed_table;
ROLLBACK;
BEGIN;
-- Since we are in a transaction, the copy should be locally executed.
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
ROLLBACK;
-- Since we are not in a transaction, the copy should not be locally executed.
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
BEGIN;
-- Since we are in a transaction, the copy should be locally executed. But
-- we are putting duplicate key, so it should error.
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 100
2, 200
3, 300
4, 400
5, 500
\.
ROLLBACK;
TRUNCATE distributed_table;
COPY distributed_table FROM STDIN WITH delimiter ',';
1, 9
\.
BEGIN;
-- Since we are in a transaction, the execution will be local, however we are putting invalid age.
-- The constaints should give an error
COPY distributed_table FROM STDIN WITH delimiter ',';
1,9
\.
ROLLBACK;
TRUNCATE distributed_table;
-- different delimiters
BEGIN;
-- run select with local execution
SELECT count(*) FROM distributed_table WHERE key = 1;
-- initial size
SELECT count(*) FROM distributed_table;
COPY distributed_table FROM STDIN WITH delimiter '|';
1|10
2|30
3|40
\.
-- new size
SELECT count(*) FROM distributed_table;
ROLLBACK;
BEGIN;
-- run select with local execution
SELECT count(*) FROM distributed_table WHERE key = 1;
-- initial size
SELECT count(*) FROM distributed_table;
COPY distributed_table FROM STDIN WITH delimiter '[';
1[10
2[30
3[40
\.
-- new size
SELECT count(*) FROM distributed_table;
ROLLBACK;
-- multiple local copies
BEGIN;
COPY distributed_table FROM STDIN WITH delimiter ',';
1,15
2,20
3,30
\.
COPY distributed_table FROM STDIN WITH delimiter ',';
10,15
20,20
30,30
\.
COPY distributed_table FROM STDIN WITH delimiter ',';
100,15
200,20
300,30
\.
ROLLBACK;
-- local copy followed by local copy should see the changes
-- and error since it is a duplicate primary key.
BEGIN;
COPY distributed_table FROM STDIN WITH delimiter ',';
1,15
\.
COPY distributed_table FROM STDIN WITH delimiter ',';
1,16
\.
ROLLBACK;
-- local copy followed by local copy should see the changes
BEGIN;
COPY distributed_table FROM STDIN WITH delimiter ',';
1,15
\.
-- select should see the change
SELECT key FROM distributed_table WHERE key = 1;
ROLLBACK;
\c - - - :master_port
SET search_path TO local_shard_copy;
SET citus.log_local_commands TO ON;
TRUNCATE TABLE reference_table;
TRUNCATE TABLE local_table;
SELECT count(*) FROM reference_table, local_table WHERE reference_table.key = local_table.key;
SET citus.enable_local_execution = 'on';
BEGIN;
-- copy should be executed locally
COPY reference_table FROM STDIN;
1
2
3
4
\.
ROLLBACK;
SET citus.enable_local_execution = 'off';
BEGIN;
-- copy should not be executed locally as citus.enable_local_execution = off
COPY reference_table FROM STDIN;
1
2
3
4
\.
ROLLBACK;
SET citus.enable_local_execution = 'on';
SET client_min_messages TO ERROR;
SET search_path TO public;
DROP SCHEMA local_shard_copy CASCADE;