Prevent downgrades when there is a single-shard table in the cluster (#6908)

Also add a few tests for Citus/PG upgrade/downgrade scenarios.

src/backend/distributed/sql/downgrades/citus--12.0-1--11.3-1.sql

@@ -1,2 +1,21 @@
 -- citus--12.0-1--11.3-1
--- this is an empty downgrade path since citus--11.3-1--12.0-1.sql is empty for now
+
+-- Throw an error if user has any distributed tables without a shard key.
+DO $$
+BEGIN
+    IF EXISTS (
+        SELECT 1 FROM pg_dist_partition
+        WHERE repmodel != 't' AND partmethod = 'n' AND colocationid != 0)
+    THEN
+        RAISE EXCEPTION 'cannot downgrade Citus because there are '
+                        'distributed tables without a shard key.'
+        USING HINT = 'You can find the distributed tables without a shard '
+                     'key in the cluster by using the following query: '
+                     '"SELECT * FROM citus_tables WHERE distribution_column '
+                     '= ''<none>'' AND colocation_id > 0".',
+        DETAIL = 'To downgrade Citus to an older version, you should '
+                 'first convert those tables to Postgres tables by '
+                 'executing SELECT undistribute_table("%s").';
+    END IF;
+END;
+$$ LANGUAGE plpgsql;

src/test/regress/after_pg_upgrade_schedule

@@ -1,4 +1,4 @@
-test: upgrade_basic_after upgrade_ref2ref_after upgrade_type_after upgrade_distributed_function_after upgrade_rebalance_strategy_after upgrade_list_citus_objects upgrade_autoconverted_after upgrade_citus_stat_activity upgrade_citus_locks
+test: upgrade_basic_after upgrade_ref2ref_after upgrade_type_after upgrade_distributed_function_after upgrade_rebalance_strategy_after upgrade_list_citus_objects upgrade_autoconverted_after upgrade_citus_stat_activity upgrade_citus_locks upgrade_single_shard_table_after
 
 # This test cannot be run with run_test.py currently due to its dependence on
 # the specific PG versions that we use to run upgrade tests. For now we leave

src/test/regress/before_pg_upgrade_schedule

@@ -5,7 +5,7 @@ test: upgrade_basic_before
 test: upgrade_ref2ref_before
 test: upgrade_type_before
 test: upgrade_distributed_function_before upgrade_rebalance_strategy_before
-test: upgrade_autoconverted_before
+test: upgrade_autoconverted_before upgrade_single_shard_table_before
 test: upgrade_citus_stat_activity
 test: upgrade_citus_locks
 test: upgrade_distributed_triggers_before

src/test/regress/expected/multi_extension.out

@@ -1332,6 +1332,21 @@ SELECT * FROM multi_extension.print_extension_changes();
 
 -- Test downgrade to 11.3-1 from 12.0-1
 ALTER EXTENSION citus UPDATE TO '12.0-1';
+CREATE TABLE null_shard_key (x int, y int);
+SELECT create_distributed_table('null_shard_key', null);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- Show that we cannot downgrade to 11.3-1 becuase the cluster has a
+-- distributed table with single-shard.
+ALTER EXTENSION citus UPDATE TO '11.3-1';
+ERROR:  cannot downgrade Citus because there are distributed tables without a shard key.
+DETAIL:  To downgrade Citus to an older version, you should first convert those tables to Postgres tables by executing SELECT undistribute_table("%s").
+HINT:  You can find the distributed tables without a shard key in the cluster by using the following query: "SELECT * FROM citus_tables WHERE distribution_column = '<none>' AND colocation_id > 0".
+CONTEXT:  PL/pgSQL function inline_code_block line XX at RAISE
+DROP TABLE null_shard_key;
 ALTER EXTENSION citus UPDATE TO '11.3-1';
 -- Should be empty result since upgrade+downgrade should be a no-op
 SELECT * FROM multi_extension.print_extension_changes();

src/test/regress/expected/upgrade_single_shard_table_after.out

@@ -0,0 +1,37 @@
+-- check that we properly retained the single-shard table
+SELECT 1 FROM pg_dist_partition
+WHERE logicalrelid = 'citus_schema.null_shard_key'::regclass AND
+      partmethod = 'n' AND repmodel = 's' AND colocationid != 0;
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+BEGIN;
+  INSERT INTO citus_schema.null_shard_key (name) VALUES ('c');
+  SELECT * FROM citus_schema.null_shard_key ORDER BY id;
+ id | name
+---------------------------------------------------------------------
+  1 | a
+  2 | b
+    | c
+(3 rows)
+
+ROLLBACK;
+-- Check that we can create a distributed table with a single-shard
+-- after upgrade.
+CREATE TABLE citus_schema.null_shard_key_after_upgrade (id bigserial, name text);
+SELECT create_distributed_table('citus_schema.null_shard_key_after_upgrade', null);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO citus_schema.null_shard_key_after_upgrade (name) VALUES ('c');
+SELECT * FROM citus_schema.null_shard_key_after_upgrade ORDER BY id;
+ id | name
+---------------------------------------------------------------------
+  1 | c
+(1 row)
+
+DROP TABLE citus_schema.null_shard_key_after_upgrade;

src/test/regress/expected/upgrade_single_shard_table_before.out

@@ -0,0 +1,8 @@
+CREATE TABLE null_shard_key (id int, name text);
+SELECT create_distributed_table('null_shard_key', null);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO null_shard_key (id, name) VALUES (1, 'a'), (2, 'b');

src/test/regress/sql/multi_extension.sql

@@ -593,6 +593,16 @@ SELECT * FROM multi_extension.print_extension_changes();
 
 -- Test downgrade to 11.3-1 from 12.0-1
 ALTER EXTENSION citus UPDATE TO '12.0-1';
+
+CREATE TABLE null_shard_key (x int, y int);
+SELECT create_distributed_table('null_shard_key', null);
+
+-- Show that we cannot downgrade to 11.3-1 becuase the cluster has a
+-- distributed table with single-shard.
+ALTER EXTENSION citus UPDATE TO '11.3-1';
+
+DROP TABLE null_shard_key;
+
 ALTER EXTENSION citus UPDATE TO '11.3-1';
 -- Should be empty result since upgrade+downgrade should be a no-op
 SELECT * FROM multi_extension.print_extension_changes();

src/test/regress/sql/upgrade_single_shard_table_after.sql

@@ -0,0 +1,18 @@
+-- check that we properly retained the single-shard table
+SELECT 1 FROM pg_dist_partition
+WHERE logicalrelid = 'citus_schema.null_shard_key'::regclass AND
+      partmethod = 'n' AND repmodel = 's' AND colocationid != 0;
+
+BEGIN;
+  INSERT INTO citus_schema.null_shard_key (name) VALUES ('c');
+  SELECT * FROM citus_schema.null_shard_key ORDER BY id;
+ROLLBACK;
+
+-- Check that we can create a distributed table with a single-shard
+-- after upgrade.
+CREATE TABLE citus_schema.null_shard_key_after_upgrade (id bigserial, name text);
+SELECT create_distributed_table('citus_schema.null_shard_key_after_upgrade', null);
+INSERT INTO citus_schema.null_shard_key_after_upgrade (name) VALUES ('c');
+SELECT * FROM citus_schema.null_shard_key_after_upgrade ORDER BY id;
+
+DROP TABLE citus_schema.null_shard_key_after_upgrade;

src/test/regress/sql/upgrade_single_shard_table_before.sql

@@ -0,0 +1,3 @@
+CREATE TABLE null_shard_key (id int, name text);
+SELECT create_distributed_table('null_shard_key', null);
+INSERT INTO null_shard_key (id, name) VALUES (1, 'a'), (2, 'b');