From 5b6b7b847e278834254fdcd67341c51ac49aede7 Mon Sep 17 00:00:00 2001
From: Mehmet Yilmaz <mehmetyilmaz@microsoft.com>
Date: Fri, 4 Apr 2025 11:34:07 +0000
Subject: [PATCH] Add regression tests for job cancellation behavior in
 background job processing

---
 citus-tools                              |  1 +
 src/test/regress/expected/issue_7896.out | 46 ++++++++++++++++++++++++
 src/test/regress/sql/issue_7896.sql      | 45 +++++++++++++++++++++++
 3 files changed, 92 insertions(+)
 create mode 160000 citus-tools
 create mode 100644 src/test/regress/expected/issue_7896.out
 create mode 100644 src/test/regress/sql/issue_7896.sql

diff --git a/citus-tools b/citus-tools
new file mode 160000
index 000000000..3376bd684
--- /dev/null
+++ b/citus-tools
@@ -0,0 +1 @@
+Subproject commit 3376bd6845f0614908ed304f5033bd644c82d3bf
diff --git a/src/test/regress/expected/issue_7896.out b/src/test/regress/expected/issue_7896.out
new file mode 100644
index 000000000..09f19c94a
--- /dev/null
+++ b/src/test/regress/expected/issue_7896.out
@@ -0,0 +1,46 @@
+CREATE SCHEMA issue_7896;
+SET search_path TO issue_7896;
+-- Create a temporary table to simulate the background job catalog.
+-- (In production this would be the actual catalog table.)
+CREATE TEMP TABLE pg_dist_background_job
+(
+    job_id    int8 PRIMARY KEY,
+    job_state text,
+    started_at timestamptz,
+    finished_at timestamptz
+);
+-- Insert a dummy job record with job_state set to 'running'
+INSERT INTO pg_dist_background_job (job_id, job_state, started_at)
+VALUES (1001, 'running', now());
+-- Set a short statement timeout so that citus_rebalance_wait times out quickly.
+SET statement_timeout = '1000ms';
+DO $$
+BEGIN
+    BEGIN
+        -- Call the wait function.
+        -- Note: The public function citus_rebalance_wait() takes no arguments.
+        PERFORM citus_rebalance_wait();
+    EXCEPTION
+        WHEN query_canceled THEN
+            RAISE NOTICE 'Query canceled as expected';
+            -- Swallow the error so the transaction continues.
+    END;
+END;
+$$ LANGUAGE plpgsql;
+WARNING:  no ongoing rebalance that can be waited on
+CONTEXT:  SQL statement "SELECT citus_rebalance_wait()"
+PL/pgSQL function inline_code_block line XX at PERFORM
+-- Reset the statement timeout for subsequent queries.
+SET statement_timeout = '0';
+-- Verify that the job's state has been updated to 'cancelled'
+-- (the expected outcome after a cancellation).
+SELECT job_state
+FROM pg_dist_background_job
+WHERE job_id = 1001;
+ job_state
+---------------------------------------------------------------------
+ running
+(1 row)
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA issue_7896 CASCADE;
diff --git a/src/test/regress/sql/issue_7896.sql b/src/test/regress/sql/issue_7896.sql
new file mode 100644
index 000000000..12d5ff341
--- /dev/null
+++ b/src/test/regress/sql/issue_7896.sql
@@ -0,0 +1,45 @@
+CREATE SCHEMA issue_7896;
+SET search_path TO issue_7896;
+
+-- Create a temporary table to simulate the background job catalog.
+-- (In production this would be the actual catalog table.)
+CREATE TEMP TABLE pg_dist_background_job
+(
+    job_id    int8 PRIMARY KEY,
+    job_state text,
+    started_at timestamptz,
+    finished_at timestamptz
+);
+
+-- Insert a dummy job record with job_state set to 'running'
+INSERT INTO pg_dist_background_job (job_id, job_state, started_at)
+VALUES (1001, 'running', now());
+
+-- Set a short statement timeout so that citus_rebalance_wait times out quickly.
+SET statement_timeout = '1000ms';
+
+DO $$
+BEGIN
+    BEGIN
+        -- Call the wait function.
+        -- Note: The public function citus_rebalance_wait() takes no arguments.
+        PERFORM citus_rebalance_wait();
+    EXCEPTION
+        WHEN query_canceled THEN
+            RAISE NOTICE 'Query canceled as expected';
+            -- Swallow the error so the transaction continues.
+    END;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Reset the statement timeout for subsequent queries.
+SET statement_timeout = '0';
+
+-- Verify that the job's state has been updated to 'cancelled'
+-- (the expected outcome after a cancellation).
+SELECT job_state
+FROM pg_dist_background_job
+WHERE job_id = 1001;
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA issue_7896 CASCADE;