From da527951ca620d8facd3f18b90cdef0be5007597 Mon Sep 17 00:00:00 2001
From: Nils Dijk <nils@citusdata.com>
Date: Tue, 13 Sep 2022 23:19:31 +0200
Subject: [PATCH] Fix: rebalance stop non super user (#6334)

No need for description, fixing issue introduced with new feature for
11.1

Fixes #6333

Due to Postgres' C api being o-indexed and postgres' attributes being
1-indexed, we were reading the wrong Datum as the Task owner when
cancelling. Here we add a test to show the error and fix the off-by-one
error.
---
 .../distributed/metadata/metadata_utility.c   |  2 +-
 .../regress/expected/background_rebalance.out | 34 +++++++++++++++++++
 src/test/regress/sql/background_rebalance.sql | 18 ++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/src/backend/distributed/metadata/metadata_utility.c b/src/backend/distributed/metadata/metadata_utility.c
index b7d20a9b5..c09ad358c 100644
--- a/src/backend/distributed/metadata/metadata_utility.c
+++ b/src/backend/distributed/metadata/metadata_utility.c
@@ -4022,7 +4022,7 @@ CancelTasksForJob(int64 jobid)
 		}
 
 		/* make sure the current user has the rights to cancel this task */
-		Oid taskOwner = DatumGetObjectId(values[Anum_pg_dist_background_task_owner]);
+		Oid taskOwner = DatumGetObjectId(values[Anum_pg_dist_background_task_owner - 1]);
 		if (superuser_arg(taskOwner) && !superuser())
 		{
 			/* must be a superuser to cancel tasks owned by superuser */
diff --git a/src/test/regress/expected/background_rebalance.out b/src/test/regress/expected/background_rebalance.out
index 32a5e86b0..8843654d6 100644
--- a/src/test/regress/expected/background_rebalance.out
+++ b/src/test/regress/expected/background_rebalance.out
@@ -176,5 +176,39 @@ SELECT citus_rebalance_wait();
 
 (1 row)
 
+DROP TABLE t1;
+-- make sure a non-super user can stop rebalancing
+CREATE USER non_super_user_rebalance WITH LOGIN;
+GRANT ALL ON SCHEMA background_rebalance TO non_super_user_rebalance;
+SET ROLE non_super_user_rebalance;
+CREATE TABLE non_super_user_t1 (a int PRIMARY KEY);
+SELECT create_distributed_table('non_super_user_t1', 'a', shard_count => 4, colocate_with => 'none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT citus_move_shard_placement(85674008, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode => 'block_writes');
+ citus_move_shard_placement
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT 1 FROM citus_rebalance_start();
+NOTICE:  Scheduled 1 moves as job xxx
+DETAIL:  Rebalance scheduled as background job
+HINT:  To monitor progress, run: SELECT * FROM pg_dist_background_task WHERE job_id = xxx ORDER BY task_id ASC; or SELECT * FROM get_rebalance_progress();
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+SELECT citus_rebalance_stop();
+ citus_rebalance_stop
+---------------------------------------------------------------------
+
+(1 row)
+
+RESET ROLE;
 SET client_min_messages TO WARNING;
 DROP SCHEMA background_rebalance CASCADE;
diff --git a/src/test/regress/sql/background_rebalance.sql b/src/test/regress/sql/background_rebalance.sql
index 9158fc532..6528c71b7 100644
--- a/src/test/regress/sql/background_rebalance.sql
+++ b/src/test/regress/sql/background_rebalance.sql
@@ -59,6 +59,24 @@ SELECT 1 FROM citus_rebalance_start();
 SELECT rebalance_table_shards();
 SELECT citus_rebalance_wait();
 
+DROP TABLE t1;
+
+
+-- make sure a non-super user can stop rebalancing
+CREATE USER non_super_user_rebalance WITH LOGIN;
+GRANT ALL ON SCHEMA background_rebalance TO non_super_user_rebalance;
+
+SET ROLE non_super_user_rebalance;
+
+CREATE TABLE non_super_user_t1 (a int PRIMARY KEY);
+SELECT create_distributed_table('non_super_user_t1', 'a', shard_count => 4, colocate_with => 'none');
+SELECT citus_move_shard_placement(85674008, 'localhost', :worker_1_port, 'localhost', :worker_2_port, shard_transfer_mode => 'block_writes');
+
+SELECT 1 FROM citus_rebalance_start();
+SELECT citus_rebalance_stop();
+
+RESET ROLE;
+
 
 SET client_min_messages TO WARNING;
 DROP SCHEMA background_rebalance CASCADE;