From aab9f623eb172acd0e90b5d064539800421a00f5 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Fri, 16 Nov 2018 23:27:34 +0100
Subject: [PATCH] Check table ownership in upgrade_to_reference_table

---
 .../distributed/utils/reference_table_utils.c        |  3 ++-
 src/test/regress/expected/multi_multiuser.out        | 12 ++++++++++++
 src/test/regress/sql/multi_multiuser.sql             |  8 ++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/utils/reference_table_utils.c b/src/backend/distributed/utils/reference_table_utils.c
index a6b0b29f0..f5b1608a2 100644
--- a/src/backend/distributed/utils/reference_table_utils.c
+++ b/src/backend/distributed/utils/reference_table_utils.c
@@ -60,8 +60,9 @@ upgrade_to_reference_table(PG_FUNCTION_ARGS)
 	uint64 shardId = INVALID_SHARD_ID;
 	DistTableCacheEntry *tableEntry = NULL;
 
-	EnsureCoordinator();
 	CheckCitusVersion(ERROR);
+	EnsureCoordinator();
+	EnsureTableOwner(relationId);
 
 	if (!IsDistributedTable(relationId))
 	{
diff --git a/src/test/regress/expected/multi_multiuser.out b/src/test/regress/expected/multi_multiuser.out
index d040f3e3a..a86ea2e44 100644
--- a/src/test/regress/expected/multi_multiuser.out
+++ b/src/test/regress/expected/multi_multiuser.out
@@ -21,6 +21,14 @@ SELECT create_distributed_table('test', 'id');
  
 (1 row)
 
+SET citus.shard_count TO 1;
+CREATE TABLE singleshard (id integer, val integer);
+SELECT create_distributed_table('singleshard', 'id');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
 -- turn off propagation to avoid Enterprise processing the following section
 SET citus.enable_ddl_propagation TO off;
 CREATE USER full_access;
@@ -212,6 +220,9 @@ ERROR:  permission denied for table test
 ABORT;
 SELECT * FROM citus_stat_statements_reset();
 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
 -- 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);
@@ -231,6 +242,7 @@ SELECT result FROM run_command_on_workers($$SELECT tableowner FROM pg_tables WHE
 
 DROP TABLE my_table;
 DROP TABLE test;
+DROP TABLE singleshard;
 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 43306f09c..478e02925 100644
--- a/src/test/regress/sql/multi_multiuser.sql
+++ b/src/test/regress/sql/multi_multiuser.sql
@@ -16,6 +16,10 @@ SET citus.shard_replication_factor TO 1;
 CREATE TABLE test (id integer, val integer);
 SELECT create_distributed_table('test', 'id');
 
+SET citus.shard_count TO 1;
+CREATE TABLE singleshard (id integer, val integer);
+SELECT create_distributed_table('singleshard', 'id');
+
 -- turn off propagation to avoid Enterprise processing the following section
 SET citus.enable_ddl_propagation TO off;
 
@@ -138,6 +142,9 @@ ABORT;
 
 SELECT * FROM citus_stat_statements_reset();
 
+-- should not be allowed to upgrade to reference table
+SELECT upgrade_to_reference_table('singleshard');
+
 -- 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);
@@ -147,6 +154,7 @@ SELECT result FROM run_command_on_workers($$SELECT tableowner FROM pg_tables WHE
 
 DROP TABLE my_table;
 DROP TABLE test;
+DROP TABLE singleshard;
 DROP USER full_access;
 DROP USER read_access;
 DROP USER no_access;