Merge pull request #4356 from citusdata/add-test-for-citus-size-func

Add test for citus table size func in transaction with modification

src/test/regress/expected/multi_cluster_management.out

@@ -309,7 +309,7 @@ SELECT create_distributed_table('cluster_management_test_colocated', 'col_1', 'h
 (1 row)
 
 -- Check that colocated shards don't get created for shards that are to be deleted
-SELECT logicalrelid, shardid, shardstate, nodename, nodeport FROM pg_dist_shard_placement NATURAL JOIN pg_dist_shard;
+SELECT logicalrelid, shardid, shardstate, nodename, nodeport FROM pg_dist_shard_placement NATURAL JOIN pg_dist_shard ORDER BY shardstate, shardid;
            logicalrelid            | shardid | shardstate | nodename  | nodeport
 ---------------------------------------------------------------------
  cluster_management_test           | 1220000 |          1 | localhost |    57637

src/test/regress/expected/single_node.out

@@ -237,6 +237,69 @@ drop cascades to table "Quoed.Schema".simple_table_name_90630518
 drop cascades to table "Quoed.Schema".simple_table_name_90630519
 drop cascades to table "Quoed.Schema".simple_table_name_90630520
 drop cascades to table "Quoed.Schema".simple_table_name_90630521
+-- test citus size functions in transaction with modification
+CREATE TABLE test_citus_size_func (a int);
+SELECT create_distributed_table('test_citus_size_func', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO test_citus_size_func VALUES(1), (2);
+BEGIN;
+	-- DDL with citus_table_size
+	ALTER TABLE test_citus_size_func ADD COLUMN newcol INT;
+	SELECT citus_table_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+	-- DDL with citus_relation_size
+	ALTER TABLE test_citus_size_func ADD COLUMN newcol INT;
+	SELECT citus_relation_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+	-- DDL with citus_total_relation_size
+	ALTER TABLE test_citus_size_func ADD COLUMN newcol INT;
+	SELECT citus_total_relation_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+    -- single shard insert with citus_table_size
+	INSERT INTO test_citus_size_func VALUES (3);
+	SELECT citus_table_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+    -- multi shard modification with citus_table_size
+	INSERT INTO test_citus_size_func  SELECT * FROM  test_citus_size_func;
+	SELECT citus_table_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+    -- single shard insert with citus_relation_size
+	INSERT INTO test_citus_size_func VALUES (3);
+	SELECT citus_relation_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+    -- multi shard modification with citus_relation_size
+	INSERT INTO test_citus_size_func  SELECT * FROM  test_citus_size_func;
+	SELECT citus_relation_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+    -- single shard insert with citus_total_relation_size
+	INSERT INTO test_citus_size_func VALUES (3);
+	SELECT citus_total_relation_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
+BEGIN;
+    -- multi shard modification with citus_total_relation_size
+	INSERT INTO test_citus_size_func  SELECT * FROM  test_citus_size_func;
+	SELECT citus_total_relation_size('test_citus_size_func');
+ERROR:  citus size functions cannot be called in transaction blocks which contain multi-shard data modifications
+ROLLBACK;
 -- we should be able to limit intermediate results
 BEGIN;
        SET LOCAL citus.max_intermediate_result_size TO 0;

src/test/regress/sql/multi_cluster_management.sql

@@ -123,7 +123,7 @@ CREATE TABLE cluster_management_test_colocated (col_1 text, col_2 int);
 SELECT create_distributed_table('cluster_management_test_colocated', 'col_1', 'hash', colocate_with=>'cluster_management_test');
 
 -- Check that colocated shards don't get created for shards that are to be deleted
-SELECT logicalrelid, shardid, shardstate, nodename, nodeport FROM pg_dist_shard_placement NATURAL JOIN pg_dist_shard;
+SELECT logicalrelid, shardid, shardstate, nodename, nodeport FROM pg_dist_shard_placement NATURAL JOIN pg_dist_shard ORDER BY shardstate, shardid;
 
 -- try to remove a node with only to be deleted placements and see that removal still fails
 SELECT master_remove_node('localhost', :worker_2_port);

src/test/regress/sql/single_node.sql

@@ -121,6 +121,65 @@ ALTER TABLE simple_table_name RENAME CONSTRAINT "looo oooo ooooo ooooooooooooooo
 SET search_path TO single_node;
 DROP SCHEMA  "Quoed.Schema" CASCADE;
 
+-- test citus size functions in transaction with modification
+CREATE TABLE test_citus_size_func (a int);
+SELECT create_distributed_table('test_citus_size_func', 'a');
+INSERT INTO test_citus_size_func VALUES(1), (2);
+
+BEGIN;
+	-- DDL with citus_table_size
+	ALTER TABLE test_citus_size_func ADD COLUMN newcol INT;
+	SELECT citus_table_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- DDL with citus_relation_size
+	ALTER TABLE test_citus_size_func ADD COLUMN newcol INT;
+	SELECT citus_relation_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- DDL with citus_total_relation_size
+	ALTER TABLE test_citus_size_func ADD COLUMN newcol INT;
+	SELECT citus_total_relation_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- single shard insert with citus_table_size
+	INSERT INTO test_citus_size_func VALUES (3);
+	SELECT citus_table_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- multi shard modification with citus_table_size
+	INSERT INTO test_citus_size_func  SELECT * FROM  test_citus_size_func;
+	SELECT citus_table_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- single shard insert with citus_relation_size
+	INSERT INTO test_citus_size_func VALUES (3);
+	SELECT citus_relation_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- multi shard modification with citus_relation_size
+	INSERT INTO test_citus_size_func  SELECT * FROM  test_citus_size_func;
+	SELECT citus_relation_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- single shard insert with citus_total_relation_size
+	INSERT INTO test_citus_size_func VALUES (3);
+	SELECT citus_total_relation_size('test_citus_size_func');
+ROLLBACK;
+
+BEGIN;
+	-- multi shard modification with citus_total_relation_size
+	INSERT INTO test_citus_size_func  SELECT * FROM  test_citus_size_func;
+	SELECT citus_total_relation_size('test_citus_size_func');
+ROLLBACK;
+
 -- we should be able to limit intermediate results
 BEGIN;
        SET LOCAL citus.max_intermediate_result_size TO 0;