Add udf tests with citus local tables (#4154)

src/test/regress/expected/citus_local_tables.out

@@ -485,10 +485,175 @@ NOTICE:  executing the command locally: TRUNCATE TABLE citus_local_tables_test_s
 -- test vacuum
 VACUUM citus_local_table_1;
 VACUUM citus_local_table_1, distributed_table, local_table, reference_table;
+-- test drop
 DROP TABLE citus_local_table_1, citus_local_table_2, distributed_table, local_table, reference_table;
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS citus_local_tables_test_schema.reference_table_xxxxx CASCADE
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS citus_local_tables_test_schema.citus_local_table_2_xxxxx CASCADE
 NOTICE:  executing the command locally: DROP TABLE IF EXISTS citus_local_tables_test_schema.citus_local_table_1_xxxxx CASCADE
+-- test some other udf's with citus local tables
+CREATE TABLE citus_local_table_4(a int);
+SELECT create_citus_local_table('citus_local_table_4');
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- should work --
+-- insert some data & create an index for table size udf's
+INSERT INTO citus_local_table_4 VALUES (1), (2), (3);
+NOTICE:  executing the command locally: INSERT INTO citus_local_tables_test_schema.citus_local_table_4_1504037 AS citus_table_alias (a) VALUES (1), (2), (3)
+CREATE INDEX citus_local_table_4_idx ON citus_local_table_4(a);
+NOTICE:  executing the command locally: CREATE  INDEX   citus_local_table_4_idx_1504037 ON citus_local_tables_test_schema.citus_local_table_4_1504037 USING btree (a )
+SELECT citus_table_size('citus_local_table_4');
+ citus_table_size
+---------------------------------------------------------------------
+             8192
+(1 row)
+
+SELECT citus_total_relation_size('citus_local_table_4');
+ citus_total_relation_size
+---------------------------------------------------------------------
+                     24576
+(1 row)
+
+SELECT citus_relation_size('citus_local_table_4');
+ citus_relation_size
+---------------------------------------------------------------------
+                8192
+(1 row)
+
+BEGIN;
+  SELECT lock_relation_if_exists('citus_local_table_4', 'ACCESS SHARE');
+ lock_relation_if_exists
+---------------------------------------------------------------------
+ t
+(1 row)
+
+  SELECT count(*) FROM pg_locks where relation='citus_local_table_4'::regclass;
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
+COMMIT;
+-- hide first column (relationId) as it might change
+SELECT part_storage_type, part_method, part_key, part_replica_count, part_max_size, part_placement_policy FROM master_get_table_metadata('citus_local_table_4');
+ part_storage_type | part_method | part_key | part_replica_count | part_max_size | part_placement_policy
+---------------------------------------------------------------------
+ t                 | n           |          |                  1 |       1536000 |                     2
+(1 row)
+
+SELECT master_get_table_ddl_events('citus_local_table_4');
+                                        master_get_table_ddl_events
+---------------------------------------------------------------------
+ CREATE TABLE citus_local_tables_test_schema.citus_local_table_4 (a integer)
+ ALTER TABLE citus_local_tables_test_schema.citus_local_table_4 OWNER TO postgres
+ CREATE INDEX citus_local_table_4_idx ON citus_local_tables_test_schema.citus_local_table_4 USING btree (a)
+(3 rows)
+
+SELECT column_to_column_name(logicalrelid, partkey)
+FROM pg_dist_partition WHERE logicalrelid = 'citus_local_table_4'::regclass;
+ column_to_column_name
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT column_name_to_column('citus_local_table_4', 'a');
+                                                  column_name_to_column
+---------------------------------------------------------------------
+ {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varlevelsup 0 :varnoold 1 :varoattno 1 :location -1}
+(1 row)
+
+SELECT master_update_shard_statistics(shardid)
+FROM (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='citus_local_table_4'::regclass) as shardid;
+ master_update_shard_statistics
+---------------------------------------------------------------------
+                           8192
+(1 row)
+
+-- will always be no-op as we create the shell table from scratch
+-- while creating a citus local table, but let's see it works
+SELECT truncate_local_data_after_distributing_table('citus_local_table_4');
+ truncate_local_data_after_distributing_table
+---------------------------------------------------------------------
+
+(1 row)
+
+BEGIN;
+  SELECT worker_drop_distributed_table('citus_local_table_4');
+ worker_drop_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- should only see shard relation
+  SELECT tableName FROM pg_catalog.pg_tables WHERE tablename LIKE 'citus_local_table_4%';
+          tablename
+---------------------------------------------------------------------
+ citus_local_table_4_1504037
+(1 row)
+
+ROLLBACK;
+-- should return a single element array that only includes its own shard id
+SELECT shardid, get_colocated_shard_array(shardid)
+FROM (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='citus_local_table_4'::regclass) as shardid;
+ shardid | get_colocated_shard_array
+---------------------------------------------------------------------
+ 1504037 | {1504037}
+(1 row)
+
+BEGIN;
+  SELECT master_remove_partition_metadata('citus_local_table_4'::regclass::oid, 'citus_local_tables_test_schema', 'citus_local_table_4');
+ master_remove_partition_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- should print 0
+  select count(*) from pg_dist_partition where logicalrelid='citus_local_table_4'::regclass;
+ count
+---------------------------------------------------------------------
+     0
+(1 row)
+
+ROLLBACK;
+-- should fail --
+SELECT upgrade_to_reference_table('citus_local_table_4');
+ERROR:  cannot upgrade to reference table
+SELECT update_distributed_table_colocation('citus_local_table_4', colocate_with => 'none');
+ERROR:  relation citus_local_table_4 should be a hash distributed table
+SELECT master_create_worker_shards('citus_local_table_4', 10, 1);
+ERROR:  unsupported table partition type: n
+SELECT master_create_empty_shard('citus_local_table_4');
+ERROR:  relation "citus_local_table_4" is a citus local table
+SELECT master_apply_delete_command('DELETE FROM citus_local_table_4');
+ERROR:  cannot delete from table
+CREATE TABLE postgres_local_table (a int);
+SELECT master_append_table_to_shard(shardId, 'postgres_local_table', 'localhost', :master_port)
+FROM (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='citus_local_table_4'::regclass) as shardid;
+ERROR:  cannot append to shardId 1504037
+-- return true
+SELECT citus_table_is_visible('citus_local_table_4'::regclass::oid);
+ citus_table_is_visible
+---------------------------------------------------------------------
+ t
+(1 row)
+
+-- return false
+SELECT relation_is_a_known_shard('citus_local_table_4');
+ relation_is_a_known_shard
+---------------------------------------------------------------------
+ f
+(1 row)
+
+-- return | false | true |
+SELECT citus_table_is_visible(tableName::regclass::oid), relation_is_a_known_shard(tableName::regclass)
+FROM (SELECT tableName FROM pg_catalog.pg_tables WHERE tablename LIKE 'citus_local_table_4_%') as tableName;
+ citus_table_is_visible | relation_is_a_known_shard
+---------------------------------------------------------------------
+ f                      | t
+(1 row)
+
 -- cleanup at exit
 DROP SCHEMA citus_local_tables_test_schema, "CiTUS!LocalTables" CASCADE;
-NOTICE:  drop cascades to 15 other objects
+NOTICE:  drop cascades to 18 other objects

src/test/regress/expected/citus_local_tables_mx.out

@@ -188,7 +188,26 @@ SELECT * FROM citus_local_table_3;
 -- finally show that we do not allow defining foreign key in mx nodes
 ALTER TABLE citus_local_table_3 ADD CONSTRAINT fkey_local_to_local_2 FOREIGN KEY(l1) REFERENCES citus_local_table_4(l1);
 ERROR:  operation is not allowed on this node
--- cleanup at exit
 \c - - - :master_port
+SET search_path TO citus_local_tables_mx;
+SELECT master_remove_distributed_table_metadata_from_workers('citus_local_table_4'::regclass::oid, 'citus_local_tables_mx', 'citus_local_table_4');
+ master_remove_distributed_table_metadata_from_workers
+---------------------------------------------------------------------
+
+(1 row)
+
+-- both workers should print 0 as master_remove_distributed_table_metadata_from_workers
+-- drops the table as well
+SELECT run_command_on_workers(
+$$
+SELECT count(*) FROM pg_catalog.pg_tables WHERE tablename='citus_local_table_4'
+$$);
+ run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,0)
+ (localhost,57638,t,0)
+(2 rows)
+
+-- cleanup at exit
 DROP SCHEMA citus_local_tables_mx CASCADE;
 NOTICE:  drop cascades to 17 other objects

src/test/regress/sql/citus_local_tables.sql

@@ -362,7 +362,85 @@ TRUNCATE citus_local_table_1, citus_local_table_2, distributed_table, local_tabl
 VACUUM citus_local_table_1;
 VACUUM citus_local_table_1, distributed_table, local_table, reference_table;
 
+-- test drop
 DROP TABLE citus_local_table_1, citus_local_table_2, distributed_table, local_table, reference_table;
 
+-- test some other udf's with citus local tables
+
+CREATE TABLE citus_local_table_4(a int);
+SELECT create_citus_local_table('citus_local_table_4');
+
+-- should work --
+
+-- insert some data & create an index for table size udf's
+INSERT INTO citus_local_table_4 VALUES (1), (2), (3);
+CREATE INDEX citus_local_table_4_idx ON citus_local_table_4(a);
+
+SELECT citus_table_size('citus_local_table_4');
+SELECT citus_total_relation_size('citus_local_table_4');
+SELECT citus_relation_size('citus_local_table_4');
+
+BEGIN;
+  SELECT lock_relation_if_exists('citus_local_table_4', 'ACCESS SHARE');
+  SELECT count(*) FROM pg_locks where relation='citus_local_table_4'::regclass;
+COMMIT;
+
+-- hide first column (relationId) as it might change
+SELECT part_storage_type, part_method, part_key, part_replica_count, part_max_size, part_placement_policy FROM master_get_table_metadata('citus_local_table_4');
+SELECT master_get_table_ddl_events('citus_local_table_4');
+
+SELECT column_to_column_name(logicalrelid, partkey)
+FROM pg_dist_partition WHERE logicalrelid = 'citus_local_table_4'::regclass;
+
+SELECT column_name_to_column('citus_local_table_4', 'a');
+
+SELECT master_update_shard_statistics(shardid)
+FROM (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='citus_local_table_4'::regclass) as shardid;
+
+-- will always be no-op as we create the shell table from scratch
+-- while creating a citus local table, but let's see it works
+SELECT truncate_local_data_after_distributing_table('citus_local_table_4');
+
+BEGIN;
+  SELECT worker_drop_distributed_table('citus_local_table_4');
+
+  -- should only see shard relation
+  SELECT tableName FROM pg_catalog.pg_tables WHERE tablename LIKE 'citus_local_table_4%';
+ROLLBACK;
+
+-- should return a single element array that only includes its own shard id
+SELECT shardid, get_colocated_shard_array(shardid)
+FROM (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='citus_local_table_4'::regclass) as shardid;
+
+BEGIN;
+  SELECT master_remove_partition_metadata('citus_local_table_4'::regclass::oid, 'citus_local_tables_test_schema', 'citus_local_table_4');
+
+  -- should print 0
+  select count(*) from pg_dist_partition where logicalrelid='citus_local_table_4'::regclass;
+ROLLBACK;
+
+-- should fail --
+
+SELECT upgrade_to_reference_table('citus_local_table_4');
+SELECT update_distributed_table_colocation('citus_local_table_4', colocate_with => 'none');
+
+SELECT master_create_worker_shards('citus_local_table_4', 10, 1);
+SELECT master_create_empty_shard('citus_local_table_4');
+SELECT master_apply_delete_command('DELETE FROM citus_local_table_4');
+
+CREATE TABLE postgres_local_table (a int);
+SELECT master_append_table_to_shard(shardId, 'postgres_local_table', 'localhost', :master_port)
+FROM (SELECT shardid FROM pg_dist_shard WHERE logicalrelid='citus_local_table_4'::regclass) as shardid;
+
+-- return true
+SELECT citus_table_is_visible('citus_local_table_4'::regclass::oid);
+
+-- return false
+SELECT relation_is_a_known_shard('citus_local_table_4');
+
+-- return | false | true |
+SELECT citus_table_is_visible(tableName::regclass::oid), relation_is_a_known_shard(tableName::regclass)
+FROM (SELECT tableName FROM pg_catalog.pg_tables WHERE tablename LIKE 'citus_local_table_4_%') as tableName;
+
 -- cleanup at exit
 DROP SCHEMA citus_local_tables_test_schema, "CiTUS!LocalTables" CASCADE;

src/test/regress/sql/citus_local_tables_mx.sql

@@ -143,6 +143,17 @@ SELECT * FROM citus_local_table_3;
 -- finally show that we do not allow defining foreign key in mx nodes
 ALTER TABLE citus_local_table_3 ADD CONSTRAINT fkey_local_to_local_2 FOREIGN KEY(l1) REFERENCES citus_local_table_4(l1);
 
--- cleanup at exit
 \c - - - :master_port
+SET search_path TO citus_local_tables_mx;
+
+SELECT master_remove_distributed_table_metadata_from_workers('citus_local_table_4'::regclass::oid, 'citus_local_tables_mx', 'citus_local_table_4');
+
+-- both workers should print 0 as master_remove_distributed_table_metadata_from_workers
+-- drops the table as well
+SELECT run_command_on_workers(
+$$
+SELECT count(*) FROM pg_catalog.pg_tables WHERE tablename='citus_local_table_4'
+$$);
+
+-- cleanup at exit
 DROP SCHEMA citus_local_tables_mx CASCADE;