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.

src/backend/distributed/sql/downgrades/citus--9.5-1--9.4-1.sql

@@ -2,7 +2,27 @@
 
 SET search_path = 'pg_catalog';
 
+-- 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
 

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();

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();