From e17025e1d4484c2c16fbe4d1149792bb2062c15e Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Sat, 17 Nov 2018 00:01:14 +0100
Subject: [PATCH] Check table ownership in mark_tables_colocated

---
 src/backend/distributed/utils/colocation_utils.c |  6 +++++-
 src/test/regress/expected/multi_multiuser.out    | 14 +++++++++++---
 src/test/regress/sql/multi_multiuser.sql         | 10 +++++++---
 3 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/backend/distributed/utils/colocation_utils.c b/src/backend/distributed/utils/colocation_utils.c
index d7c0e9a2a..ff273a99c 100644
--- a/src/backend/distributed/utils/colocation_utils.c
+++ b/src/backend/distributed/utils/colocation_utils.c
@@ -78,8 +78,9 @@ mark_tables_colocated(PG_FUNCTION_ARGS)
 							   "operation")));
 	}
 
-	EnsureCoordinator();
 	CheckCitusVersion(ERROR);
+	EnsureCoordinator();
+	EnsureTableOwner(sourceRelationId);
 
 	relationIdDatumArray = DeconstructArrayObject(relationIdArrayObject);
 
@@ -87,6 +88,9 @@ mark_tables_colocated(PG_FUNCTION_ARGS)
 	{
 		Oid nextRelationOid = DatumGetObjectId(relationIdDatumArray[relationIndex]);
 
+		/* we require that the user either owns all tables or is superuser */
+		EnsureTableOwner(nextRelationOid);
+
 		MarkTablesColocated(sourceRelationId, nextRelationOid);
 	}
 
diff --git a/src/test/regress/expected/multi_multiuser.out b/src/test/regress/expected/multi_multiuser.out
index fb94d9917..5d04ddfa1 100644
--- a/src/test/regress/expected/multi_multiuser.out
+++ b/src/test/regress/expected/multi_multiuser.out
@@ -21,6 +21,13 @@ 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');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
 SET citus.shard_count TO 1;
 CREATE TABLE singleshard (id integer, val integer);
 SELECT create_distributed_table('singleshard', 'id');
@@ -234,6 +241,9 @@ ERROR:  permission denied for function citus_stat_statements_reset
 -- should not be allowed to upgrade to reference table
 SELECT upgrade_to_reference_table('singleshard');
 ERROR:  must be owner of table singleshard
+-- should not be allowed to co-located tables
+SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+ERROR:  must be owner of table test
 -- table owner should be the same on the shards, even when distributing the table as superuser
 SET ROLE full_access;
 CREATE TABLE my_table (id integer, val integer);
@@ -251,9 +261,7 @@ SELECT result FROM run_command_on_workers($$SELECT tableowner FROM pg_tables WHE
  full_access
 (2 rows)
 
-DROP TABLE my_table;
-DROP TABLE test;
-DROP TABLE singleshard;
+DROP TABLE my_table, singleshard, test, test_coloc;
 DROP USER full_access;
 DROP USER read_access;
 DROP USER no_access;
diff --git a/src/test/regress/sql/multi_multiuser.sql b/src/test/regress/sql/multi_multiuser.sql
index 160a74c7f..6b9529c33 100644
--- a/src/test/regress/sql/multi_multiuser.sql
+++ b/src/test/regress/sql/multi_multiuser.sql
@@ -16,6 +16,9 @@ SET citus.shard_replication_factor TO 1;
 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');
+
 SET citus.shard_count TO 1;
 CREATE TABLE singleshard (id integer, val integer);
 SELECT create_distributed_table('singleshard', 'id');
@@ -151,6 +154,9 @@ SELECT * FROM citus_stat_statements_reset();
 -- should not be allowed to upgrade to reference table
 SELECT upgrade_to_reference_table('singleshard');
 
+-- should not be allowed to co-located tables
+SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+
 -- table owner should be the same on the shards, even when distributing the table as superuser
 SET ROLE full_access;
 CREATE TABLE my_table (id integer, val integer);
@@ -158,9 +164,7 @@ RESET ROLE;
 SELECT create_distributed_table('my_table', 'id');
 SELECT result FROM run_command_on_workers($$SELECT tableowner FROM pg_tables WHERE tablename LIKE 'my_table_%' LIMIT 1$$);
 
-DROP TABLE my_table;
-DROP TABLE test;
-DROP TABLE singleshard;
+DROP TABLE my_table, singleshard, test, test_coloc;
 DROP USER full_access;
 DROP USER read_access;
 DROP USER no_access;