add simple test for background job

src/backend/distributed/sql/udfs/citus_wait_for_rebalance_job/11.1-1.sql

@@ -2,7 +2,7 @@ CREATE FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     RETURNS VOID
     LANGUAGE C STRICT
     AS 'MODULE_PATHNAME',$$citus_wait_for_rebalance_job$$;
-COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job()
+COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     IS 'blocks till the job identified by jobid is at a terminal state, errors if no such job exists';
 
-GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job() TO PUBLIC;
+GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint) TO PUBLIC;

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

@@ -2,7 +2,7 @@ CREATE FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     RETURNS VOID
     LANGUAGE C STRICT
     AS 'MODULE_PATHNAME',$$citus_wait_for_rebalance_job$$;
-COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job()
+COMMENT ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint)
     IS 'blocks till the job identified by jobid is at a terminal state, errors if no such job exists';
 
-GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job() TO PUBLIC;
+GRANT EXECUTE ON FUNCTION pg_catalog.citus_wait_for_rebalance_job(jobid bigint) TO PUBLIC;

src/test/regress/expected/rebalance_job.out

@@ -0,0 +1,27 @@
+CREATE SCHEMA rebalance_job;
+SET search_path TO rebalance_job;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 3536400;
+CREATE TABLE results (a int);
+-- simple job that inserts 1 into results to show that query runs
+SELECT a FROM results WHERE a = 1; -- verify result is not in there
+ a
+---------------------------------------------------------------------
+(0 rows)
+
+INSERT INTO pg_dist_rebalance_jobs (command) VALUES ($job$ INSERT INTO rebalance_job.results VALUES ( 1 ); $job$) RETURNING jobid \gset
+SELECT citus_wait_for_rebalance_job(:jobid); -- wait for the job to be finished
+ citus_wait_for_rebalance_job
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT a FROM results WHERE a = 1; -- verify result is there
+ a
+---------------------------------------------------------------------
+ 1
+(1 row)
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA rebalance_job CASCADE;

src/test/regress/multi_schedule

@@ -97,6 +97,7 @@ test: issue_5248 issue_5099
 test: object_propagation_debug
 test: undistribute_table
 test: run_command_on_all_nodes
+test: rebalance_job
 
 
 # ---------

src/test/regress/sql/rebalance_job.sql

@@ -0,0 +1,16 @@
+CREATE SCHEMA rebalance_job;
+SET search_path TO rebalance_job;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+SET citus.next_shard_id TO 3536400;
+
+CREATE TABLE results (a int);
+
+-- simple job that inserts 1 into results to show that query runs
+SELECT a FROM results WHERE a = 1; -- verify result is not in there
+INSERT INTO pg_dist_rebalance_jobs (command) VALUES ($job$ INSERT INTO rebalance_job.results VALUES ( 1 ); $job$) RETURNING jobid \gset
+SELECT citus_wait_for_rebalance_job(:jobid); -- wait for the job to be finished
+SELECT a FROM results WHERE a = 1; -- verify result is there
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA rebalance_job CASCADE;