Enhance pg upgrade tests (#3002)

* Enhance pg upgrade tests

* Add a specific upgrade test for pg_dist_partition

We store the index of distribution column, and when a column with an
index that is smaller than distribution column index is dropped before
an upgrade, the index should still match the distribution column after
an upgrade

src/test/regress/expected/upgrade_distributed_table_after.out

@@ -15,5 +15,59 @@ SELECT * FROM t WHERE a = 1;
  1
 (1 row)
 
+INSERT INTO t SELECT * FROM generate_series(10, 15);
+SELECT * FROM t WHERE a = 10;
+ a  
+----
+ 10
+(1 row)
+
+SELECT * FROM t WHERE a = 11;
+ a  
+----
+ 11
+(1 row)
+
+-- test distributed type
+INSERT INTO t1 VALUES (1, (2,3)::tc1);
+SELECT * FROM t1;
+ a |   b   
+---+-------
+ 1 | (2,3)
+(1 row)
+
+ALTER TYPE tc1 RENAME TO tc1_newname;
+INSERT INTO t1 VALUES (3, (4,5)::tc1_newname);
+TRUNCATE TABLE t;
+SELECT * FROM T;
+ a 
+---
+(0 rows)
+
+-- verify that the table whose column is dropped before a pg_upgrade still works as expected.
+SELECT * FROM t_ab ORDER BY b;
+ b  
+----
+ 11
+ 22
+ 33
+(3 rows)
+
+SELECT * FROM t_ab WHERE b = 11;
+ b  
+----
+ 11
+(1 row)
+
+SELECT * FROM t_ab WHERE b = 22;
+ b  
+----
+ 22
+(1 row)
+
+DROP TABLE t;
 DROP SCHEMA upgrade_distributed_table_before CASCADE;
-NOTICE:  drop cascades to table t
+NOTICE:  drop cascades to 3 other objects
+DETAIL:  drop cascades to type tc1_newname
+drop cascades to table t1
+drop cascades to table t_ab

src/test/regress/expected/upgrade_distributed_table_before.out

@@ -8,3 +8,25 @@ SELECT create_distributed_table('t', 'a');
 (1 row)
 
 INSERT INTO t SELECT * FROM generate_series(1, 5);
+CREATE TYPE tc1 AS (a int, b int);
+CREATE TABLE t1 (a int PRIMARY KEY, b tc1);
+SELECT create_distributed_table('t1','a');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+-- We store the index of distribution column and here we check that the distribution
+-- column index does not change after an upgrade if we drop a column that comes before the
+-- distribution column. The index information is in partkey column of pg_dist_partition table.
+CREATE TABLE t_ab(a int, b int);
+SELECT create_distributed_table('t_ab', 'b');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+INSERT INTO t_ab VALUES (1, 11);
+INSERT INTO t_ab VALUES (2, 22);
+INSERT INTO t_ab VALUES (3, 33);
+ALTER TABLE t_ab DROP a;

src/test/regress/sql/upgrade_distributed_table_after.sql

@@ -3,4 +3,26 @@ SET search_path TO upgrade_distributed_table_before, public;
 SELECT * FROM t ORDER BY a;
 SELECT * FROM t WHERE a = 1;
 
+INSERT INTO t SELECT * FROM generate_series(10, 15);
+
+SELECT * FROM t WHERE a = 10;
+SELECT * FROM t WHERE a = 11;
+
+-- test distributed type
+INSERT INTO t1 VALUES (1, (2,3)::tc1);
+SELECT * FROM t1;
+ALTER TYPE tc1 RENAME TO tc1_newname;
+INSERT INTO t1 VALUES (3, (4,5)::tc1_newname);
+
+TRUNCATE TABLE t;
+
+SELECT * FROM T;
+
+-- verify that the table whose column is dropped before a pg_upgrade still works as expected.
+SELECT * FROM t_ab ORDER BY b;
+SELECT * FROM t_ab WHERE b = 11;
+SELECT * FROM t_ab WHERE b = 22;
+
+DROP TABLE t;
+
 DROP SCHEMA upgrade_distributed_table_before CASCADE;

src/test/regress/sql/upgrade_distributed_table_before.sql

@@ -4,3 +4,19 @@ SET search_path TO upgrade_distributed_table_before, public;
 CREATE TABLE t(a int);
 SELECT create_distributed_table('t', 'a');
 INSERT INTO t SELECT * FROM generate_series(1, 5);
+
+CREATE TYPE tc1 AS (a int, b int);
+CREATE TABLE t1 (a int PRIMARY KEY, b tc1);
+SELECT create_distributed_table('t1','a');
+
+-- We store the index of distribution column and here we check that the distribution
+-- column index does not change after an upgrade if we drop a column that comes before the
+-- distribution column. The index information is in partkey column of pg_dist_partition table.
+CREATE TABLE t_ab(a int, b int);
+SELECT create_distributed_table('t_ab', 'b');
+INSERT INTO t_ab VALUES (1, 11);
+INSERT INTO t_ab VALUES (2, 22);
+INSERT INTO t_ab VALUES (3, 33);
+
+ALTER TABLE t_ab DROP a;
+