Revoke usage from the citus schema from public (#3123)

Revoke usage from the citus schema from public

src/backend/distributed/sql/citus--9.0-1--9.0-2.sql

@@ -0,0 +1,10 @@
+-- Using the citus schema is a bad idea since many environments use "citus"
+-- as the main user and the "citus" schema then sits in front of the
+-- search_path.
+REVOKE USAGE ON SCHEMA citus FROM public;
+
+-- redefine distributed_tables_colocated to avoid using citus schema
+#include "udfs/distributed_tables_colocated/9.0-2.sql"
+
+-- type was used in old version of distributed_tables_colocated
+DROP TYPE citus.colocation_placement_type;

src/backend/distributed/sql/udfs/distributed_tables_colocated/9.0-2.sql

@@ -0,0 +1,22 @@
+--
+-- distributed_tables_colocated returns true if given tables are co-located, false otherwise.
+-- The function checks shard definitions, matches shard placements for given tables.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.distributed_tables_colocated(table1 regclass,
+																   table2 regclass)
+RETURNS bool
+LANGUAGE plpgsql
+AS $function$
+DECLARE
+	table1_colocationid int;
+	table2_colocationid int;
+BEGIN
+	SELECT colocationid INTO table1_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table1;
+
+	SELECT colocationid INTO table2_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table2;
+
+	RETURN table1_colocationid = table2_colocationid;
+END;
+$function$;

src/backend/distributed/sql/udfs/distributed_tables_colocated/latest.sql

@@ -0,0 +1,22 @@
+--
+-- distributed_tables_colocated returns true if given tables are co-located, false otherwise.
+-- The function checks shard definitions, matches shard placements for given tables.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.distributed_tables_colocated(table1 regclass,
+																   table2 regclass)
+RETURNS bool
+LANGUAGE plpgsql
+AS $function$
+DECLARE
+	table1_colocationid int;
+	table2_colocationid int;
+BEGIN
+	SELECT colocationid INTO table1_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table1;
+
+	SELECT colocationid INTO table2_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table2;
+
+	RETURN table1_colocationid = table2_colocationid;
+END;
+$function$;

src/test/regress/expected/multi_citus_tools.out

@@ -421,8 +421,8 @@ SELECT create_distributed_table('check_colocated', 'key', 'hash');
  
 (1 row)
 
-CREATE TABLE second_table (key int);
 SET citus.shard_count TO 4;
+CREATE TABLE second_table (key int);
 SELECT create_distributed_table('second_table', 'key', 'hash');
  create_distributed_table 
 --------------------------
@@ -433,15 +433,13 @@ SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_tab
 												  'select 1');
 ERROR:  tables check_colocated and second_table are not co-located
 -- even when the difference is in replication factor, an error is thrown
-SELECT master_drop_all_shards('second_table'::regclass, current_schema(), 'second_table');
- master_drop_all_shards 
-------------------------
-                      4
-(1 row)
-
-SELECT master_create_worker_shards('second_table', 5, 1);
- master_create_worker_shards 
------------------------------
+DROP TABLE second_table;
+SET citus.shard_replication_factor TO 1;
+SET citus.shard_count TO 5;
+CREATE TABLE second_table (key int);
+SELECT create_distributed_table('second_table', 'key', 'hash');
+ create_distributed_table 
+--------------------------
  
 (1 row)
 
@@ -449,15 +447,13 @@ SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_tab
 												  'select 1');
 ERROR:  tables check_colocated and second_table are not co-located
 -- when everything matches, the command is run!
-SELECT master_drop_all_shards('second_table'::regclass, current_schema(), 'second_table');
- master_drop_all_shards 
-------------------------
-                      5
-(1 row)
-
-SELECT master_create_worker_shards('second_table', 5, 2);
- master_create_worker_shards 
------------------------------
+DROP TABLE second_table;
+SET citus.shard_replication_factor TO 2;
+SET citus.shard_count TO 5;
+CREATE TABLE second_table (key int);
+SELECT create_distributed_table('second_table', 'key', 'hash');
+ create_distributed_table 
+--------------------------
  
 (1 row)
 
@@ -477,36 +473,6 @@ SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_tab
  localhost |    57638 |  1240009 |  1240023 | t       | 1
 (10 rows)
 
--- when a placement is invalid considers the tables to not be colocated
-UPDATE pg_dist_shard_placement SET shardstate = 3 WHERE shardid = (
-		SELECT shardid FROM pg_dist_shard
-		WHERE nodeport = :worker_1_port AND logicalrelid = 'second_table'::regclass
-		ORDER BY 1 ASC LIMIT 1
-);
-SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
-												  'select 1');
-ERROR:  tables check_colocated and second_table are not co-located
--- when matching placement is also invalid, considers the tables to be colocated
-UPDATE pg_dist_shard_placement SET shardstate = 3 WHERE shardid = (
-		SELECT shardid FROM pg_dist_shard
-		WHERE nodeport = :worker_1_port AND logicalrelid = 'check_colocated'::regclass
-		ORDER BY 1 ASC LIMIT 1
-);
-SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
-												  'select 1');
- nodename  | nodeport | shardid1 | shardid2 | success | result 
------------+----------+----------+----------+---------+--------
- localhost |    57638 |  1240005 |  1240019 | t       | 1
- localhost |    57637 |  1240006 |  1240020 | t       | 1
- localhost |    57638 |  1240006 |  1240020 | t       | 1
- localhost |    57637 |  1240007 |  1240021 | t       | 1
- localhost |    57638 |  1240007 |  1240021 | t       | 1
- localhost |    57637 |  1240008 |  1240022 | t       | 1
- localhost |    57638 |  1240008 |  1240022 | t       | 1
- localhost |    57637 |  1240009 |  1240023 | t       | 1
- localhost |    57638 |  1240009 |  1240023 | t       | 1
-(9 rows)
-
 DROP TABLE check_colocated CASCADE;
 DROP TABLE second_table CASCADE;
 -- runs on all shards

src/test/regress/expected/multi_multiuser.out

@@ -25,7 +25,7 @@ SELECT create_distributed_table('test', 'id');
 (1 row)
 
 CREATE TABLE test_coloc (id integer, val integer);
-SELECT create_distributed_table('test_coloc', 'id', colocate_with := 'none');
+SELECT create_distributed_table('test_coloc', 'id', colocate_with := 'test');
  create_distributed_table 
 --------------------------
  
@@ -431,7 +431,7 @@ INSERT INTO full_access_user_schema.t1 VALUES (1),(2),(3);
 -- not allowed to create a table
 SELECT create_distributed_table('full_access_user_schema.t1', 'id');
 ERROR:  permission denied for schema full_access_user_schema
-CONTEXT:  while executing command on localhost:57637
+CONTEXT:  while executing command on localhost:57638
 RESET ROLE;
 SET ROLE usage_access;
 CREATE TYPE usage_access_type AS ENUM ('a', 'b');
@@ -675,7 +675,7 @@ ERROR:  could not receive file "base/pgsql_job_cache/job_0042/task_000001/p_0000
 -- different user should not be able to fetch partition file
 SET ROLE usage_access;
 SELECT worker_fetch_partition_file(42, 1, 1, 1, 'localhost', :worker_1_port);
-WARNING:  could not open file "base/pgsql_job_cache/job_0042/task_000001/p_00001.17996": No such file or directory
+WARNING:  could not open file "base/pgsql_job_cache/job_0042/task_000001/p_00001.37455": No such file or directory
 CONTEXT:  while executing command on localhost:57637
 ERROR:  could not receive file "base/pgsql_job_cache/job_0042/task_000001/p_00001" from localhost:57637
 -- only the user whom created the files should be able to fetch
@@ -714,7 +714,7 @@ RESET ROLE;
 -- test that the super user is unable to read the contents of the intermediate file,
 -- although it does create the table
 SELECT worker_merge_files_into_table(42, 1, ARRAY['a'], ARRAY['integer']);
-WARNING:  Task file "task_000001.17982" does not have expected suffix ".10"
+WARNING:  Task file "task_000001.36164" does not have expected suffix ".10"
  worker_merge_files_into_table 
 -------------------------------
  
@@ -756,7 +756,7 @@ SELECT worker_merge_files_and_run_query(42, 1,
     'CREATE TABLE task_000001_merge(merge_column_0 int)',
     'CREATE TABLE task_000001 (a) AS SELECT sum(merge_column_0) FROM task_000001_merge'
 );
-WARNING:  Task file "task_000001.17982" does not have expected suffix ".10"
+WARNING:  Task file "task_000001.36164" does not have expected suffix ".10"
  worker_merge_files_and_run_query 
 ----------------------------------
  

src/test/regress/sql/multi_citus_tools.sql

@@ -220,38 +220,30 @@ DROP TABLE check_placements CASCADE;
 -- make sure run_on_all_colocated_placements correctly detects colocation
 CREATE TABLE check_colocated (key int);
 SELECT create_distributed_table('check_colocated', 'key', 'hash');
-CREATE TABLE second_table (key int);
 
 SET citus.shard_count TO 4;
+CREATE TABLE second_table (key int);
 SELECT create_distributed_table('second_table', 'key', 'hash');
 SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
 												  'select 1');
 -- even when the difference is in replication factor, an error is thrown
-SELECT master_drop_all_shards('second_table'::regclass, current_schema(), 'second_table');
-SELECT master_create_worker_shards('second_table', 5, 1);
+DROP TABLE second_table;
+SET citus.shard_replication_factor TO 1;
+SET citus.shard_count TO 5;
+CREATE TABLE second_table (key int);
+SELECT create_distributed_table('second_table', 'key', 'hash');
 SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
 												  'select 1');
 -- when everything matches, the command is run!
-SELECT master_drop_all_shards('second_table'::regclass, current_schema(), 'second_table');
-SELECT master_create_worker_shards('second_table', 5, 2);
-SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
-												  'select 1');
--- when a placement is invalid considers the tables to not be colocated
-UPDATE pg_dist_shard_placement SET shardstate = 3 WHERE shardid = (
-		SELECT shardid FROM pg_dist_shard
-		WHERE nodeport = :worker_1_port AND logicalrelid = 'second_table'::regclass
-		ORDER BY 1 ASC LIMIT 1
-);
-SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
-												  'select 1');
--- when matching placement is also invalid, considers the tables to be colocated
-UPDATE pg_dist_shard_placement SET shardstate = 3 WHERE shardid = (
-		SELECT shardid FROM pg_dist_shard
-		WHERE nodeport = :worker_1_port AND logicalrelid = 'check_colocated'::regclass
-		ORDER BY 1 ASC LIMIT 1
-);
+DROP TABLE second_table;
+SET citus.shard_replication_factor TO 2;
+SET citus.shard_count TO 5;
+CREATE TABLE second_table (key int);
+SELECT create_distributed_table('second_table', 'key', 'hash');
+
 SELECT * FROM run_command_on_colocated_placements('check_colocated', 'second_table',
 												  'select 1');
+
 DROP TABLE check_colocated CASCADE;
 DROP TABLE second_table CASCADE;
 

src/test/regress/sql/multi_multiuser.sql

@@ -21,7 +21,7 @@ CREATE TABLE test (id integer, val integer);
 SELECT create_distributed_table('test', 'id');
 
 CREATE TABLE test_coloc (id integer, val integer);
-SELECT create_distributed_table('test_coloc', 'id', colocate_with := 'none');
+SELECT create_distributed_table('test_coloc', 'id', colocate_with := 'test');
 
 SET citus.shard_count TO 1;
 CREATE TABLE singleshard (id integer, val integer);