From 64d5ac6a104129ab16c105e00caa436d94508def Mon Sep 17 00:00:00 2001
From: Onur Tirtir <onurcantirtir@gmail.com>
Date: Tue, 22 Sep 2020 14:19:50 +0300
Subject: [PATCH] Do not downgrade if a citus local table exists (#4174)

As the previous versions of Citus don't know how to handle citus local
tables, we should prevent downgrading from 9.5 to older versions if any
citus local tables exists.
---
 .../sql/downgrades/citus--9.5-1--9.4-1.sql    | 22 +++++++++++++++++-
 src/test/regress/expected/multi_extension.out | 23 ++++++++++++++++++-
 src/test/regress/sql/multi_extension.sql      | 12 ++++++++++
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/src/backend/distributed/sql/downgrades/citus--9.5-1--9.4-1.sql b/src/backend/distributed/sql/downgrades/citus--9.5-1--9.4-1.sql
index 5da7b5140..a2a9ef18f 100644
--- a/src/backend/distributed/sql/downgrades/citus--9.5-1--9.4-1.sql
+++ b/src/backend/distributed/sql/downgrades/citus--9.5-1--9.4-1.sql
@@ -2,7 +2,27 @@
 
 SET search_path = 'pg_catalog';
 
-DROP FUNCTION create_citus_local_table(table_name regclass);
+-- Check if user has any citus local tables.
+-- If not, DROP create_citus_local_table UDF and continue safely.
+-- Otherwise, raise an exception to stop the downgrade process.
+DO $$
+DECLARE
+    citus_local_table_count INTEGER;
+BEGIN
+    SELECT COUNT(*) INTO citus_local_table_count
+    FROM pg_dist_partition WHERE repmodel != 't' AND partmethod = 'n';
+
+    IF citus_local_table_count = 0 THEN
+        -- no citus local tables exist, can safely downgrade
+        DROP FUNCTION create_citus_local_table(table_name regclass);
+    ELSE
+        RAISE EXCEPTION 'citus local tables are introduced in Citus 9.5'
+        USING HINT = 'To downgrade Citus to an older version, you should '
+                     'first convert each citus local table to a postgres '
+                     'table by executing SELECT undistribute_table("%s")';
+    END IF;
+END;
+$$ LANGUAGE plpgsql;
 
 --  task_tracker_* functions
 
diff --git a/src/test/regress/expected/multi_extension.out b/src/test/regress/expected/multi_extension.out
index 85776f3a6..37f3b858e 100644
--- a/src/test/regress/expected/multi_extension.out
+++ b/src/test/regress/expected/multi_extension.out
@@ -408,6 +408,27 @@ SELECT * FROM print_extension_changes();
 
 -- Test downgrade to 9.4-1 from 9.5-1
 ALTER EXTENSION citus UPDATE TO '9.5-1';
+BEGIN;
+  SELECT master_add_node('localhost', :master_port, groupId=>0);
+ master_add_node
+---------------------------------------------------------------------
+               1
+(1 row)
+
+  CREATE TABLE citus_local_table (a int);
+  SELECT create_citus_local_table('citus_local_table');
+ create_citus_local_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- downgrade from 9.5-1 to 9.4-1 should fail as we have a citus local table
+  ALTER EXTENSION citus UPDATE TO '9.4-1';
+ERROR:  citus local tables are introduced in Citus 9.5
+HINT:  To downgrade Citus to an older version, you should first convert each citus local table to a postgres table by executing SELECT undistribute_table("%s")
+CONTEXT:  PL/pgSQL function inline_code_block line 11 at RAISE
+ROLLBACK;
+-- now we can downgrade as there is no citus local table
 ALTER EXTENSION citus UPDATE TO '9.4-1';
 -- Should be empty result since upgrade+downgrade should be a no-op
 SELECT * FROM print_extension_changes();
@@ -418,7 +439,7 @@ SELECT * FROM print_extension_changes();
 -- Snapshot of state at 9.5-1
 ALTER EXTENSION citus UPDATE TO '9.5-1';
 SELECT * FROM print_extension_changes();
-                           previous_object                           |            current_object
+                           previous_object                           |               current_object
 ---------------------------------------------------------------------
  function task_tracker_assign_task(bigint,integer,text)              |
  function task_tracker_cleanup_job(bigint)                           |
diff --git a/src/test/regress/sql/multi_extension.sql b/src/test/regress/sql/multi_extension.sql
index cad96a8f7..da73f0171 100644
--- a/src/test/regress/sql/multi_extension.sql
+++ b/src/test/regress/sql/multi_extension.sql
@@ -195,7 +195,19 @@ SELECT * FROM print_extension_changes();
 
 -- Test downgrade to 9.4-1 from 9.5-1
 ALTER EXTENSION citus UPDATE TO '9.5-1';
+
+BEGIN;
+  SELECT master_add_node('localhost', :master_port, groupId=>0);
+  CREATE TABLE citus_local_table (a int);
+  SELECT create_citus_local_table('citus_local_table');
+
+  -- downgrade from 9.5-1 to 9.4-1 should fail as we have a citus local table
+  ALTER EXTENSION citus UPDATE TO '9.4-1';
+ROLLBACK;
+
+-- now we can downgrade as there is no citus local table
 ALTER EXTENSION citus UPDATE TO '9.4-1';
+
 -- Should be empty result since upgrade+downgrade should be a no-op
 SELECT * FROM print_extension_changes();