Add Citus upgrade tests with its job (#3003)

* Add initial citus upgrade test

* Add restart databases and run tests in all nodes

* Add output for citus versions 8.0 8.1 8.2 and 8.3

* Add verify step for citus upgrade

* Add target for citus upgrade test in makefile

* Add check citus upgrade job

* Fix installation file path and add missing tar

* Run citus upgrade for v8.0 v8.1 v8.2 and v8.3

* Create upgrade_common file and rename upgrade check

* Add pg version to citus upgrade test

* Test with postgres 10 and 11 in citus upgrade tests

* Add readme for citus upgrade test

* Add some basic tests to citus upgrade tests

* Add citus upgrade mixed mode test

* Remove citus artifacts before installing another one

* Refactor citus upgrade test according to reviews

* quick and dirty rewrite of citus upgrade tests to support local execution.

I think we need to change the makefile in such a way that the tar files can be injected from the circle ci config file.

Also I removed some of the citus version checks you had to not have the requirement to pass that in separately from the pre tar file. I am not super happy with it, but two flags that need to be kept in sync is also not desirable. Instead I print out the citus version that is installed per node. This will not cause a failure if they are not what one would expect but it lets us verify we are running the expected version.

* use latest citusupgradetester in circleci

* update readme and use common alias for upgrade_common import

.circleci/config.yml

@@ -149,7 +149,7 @@ jobs:
           at: .
       - run: 
           name: 'Install and test postgres upgrade'
-          command: 'chown -R circleci:circleci /home/circleci && install-and-test-ext --target check-upgrade --old-pg-version 10 --new-pg-version 11'
+          command: 'chown -R circleci:circleci /home/circleci && install-and-test-ext --target check-pg-upgrade --old-pg-version 10 --new-pg-version 11'
 
   test-11-12_check-pg-upgrade:
     docker:
@@ -160,7 +160,7 @@ jobs:
           at: .
       - run: 
           name: 'Install and test postgres upgrade'
-          command: 'chown -R circleci:circleci /home/circleci && install-and-test-ext --target check-upgrade --old-pg-version 11 --new-pg-version 12'
+          command: 'chown -R circleci:circleci /home/circleci && install-and-test-ext --target check-pg-upgrade --old-pg-version 11 --new-pg-version 12'
 
   test-12_check-multi:
     docker:
@@ -271,6 +271,32 @@ jobs:
       - codecov/upload: 
           flags: 'test_12,isolation'        
 
+
+
+  test-11_check-citus-upgrade:
+    docker:
+      - image: 'citus/citusupgradetester-11:latest'
+    working_directory: /home/circleci/project     
+    steps:
+      - {attach_workspace: {at: .}}
+      - run: 
+          name: 'Install and test citus upgrade'
+          command: |
+           chown -R circleci:circleci /home/circleci
+           install-and-test-ext --target check-citus-upgrade --citus-pre-tar /install-pg11-citusv8.0.0.tar
+           install-and-test-ext --target check-citus-upgrade --citus-pre-tar /install-pg11-citusv8.1.0.tar 
+           install-and-test-ext --target check-citus-upgrade --citus-pre-tar /install-pg11-citusv8.2.0.tar 
+           install-and-test-ext --target check-citus-upgrade --citus-pre-tar /install-pg11-citusv8.3.0.tar 
+
+           install-and-test-ext --target check-citus-upgrade-mixed --citus-pre-tar /install-pg11-citusv8.0.0.tar
+           install-and-test-ext --target check-citus-upgrade-mixed --citus-pre-tar /install-pg11-citusv8.1.0.tar 
+           install-and-test-ext --target check-citus-upgrade-mixed --citus-pre-tar /install-pg11-citusv8.2.0.tar 
+           install-and-test-ext --target check-citus-upgrade-mixed --citus-pre-tar /install-pg11-citusv8.3.0.tar
+
+
+          
+           
+
 workflows:
   version: 2
   build_and_test:
@@ -297,10 +323,6 @@ workflows:
       - test-11_check-non-adaptive-isolation:
           requires: [build]
 
-      - test-10-11_check-pg-upgrade: 
-          requires: [build]
-      - test-11-12_check-pg-upgrade: 
-          requires: [build]    
       - test-12_check-multi:
           requires: [build] 
       - test-12_check-tt-van-mx:
@@ -319,3 +341,11 @@ workflows:
       - test-12_check-non-adaptive-isolation:
           requires: [build] 
 
+      - test-10-11_check-pg-upgrade: 
+          requires: [build]
+      - test-11-12_check-pg-upgrade: 
+          requires: [build]    
+
+      - test-11_check-citus-upgrade:
+          requires: [build]    
+             

src/test/regress/.gitignore

@@ -4,6 +4,7 @@
 # Generated subdirectories
 /tmp_check/
 /tmp_upgrade/
+/tmp_citus_upgrade/
 /results/
 /log/
 

src/test/regress/Makefile

@@ -23,7 +23,8 @@ MULTI_INSTALLDIR=$(CURDIR)/tmp_check/install
 pg_regress_multi_check = $(PERL) $(citus_abs_srcdir)/pg_regress_multi.pl --pgxsdir="$(pgxsdir)" --bindir="$(bindir)" --libdir="$(libdir)" --majorversion="$(MAJORVERSION)" --postgres-builddir="$(postgres_abs_builddir)"  --postgres-srcdir="$(postgres_abs_srcdir)"
 MULTI_REGRESS_OPTS = --inputdir=$(citus_abs_srcdir) $(pg_regress_locale_flags) --launcher="$(citus_abs_srcdir)/log_test_times"
 
-pg_upgrade_check = $(citus_abs_srcdir)/upgrade/upgrade_test.py
+pg_upgrade_check = $(citus_abs_srcdir)/upgrade/pg_upgrade_test.py
+citus_upgrade_check = $(citus_abs_srcdir)/upgrade/citus_upgrade_test.py
 
 # XXX: Can't actually do useful testruns against install - $libdir
 # etc will point to the directory configured during postgres'
@@ -142,9 +143,24 @@ check-failure-base: all
 	$(pg_regress_multi_check) --load-extension=citus --mitmproxy \
 	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/failure_base_schedule $(EXTRA_TESTS)
 
-check-upgrade:
+check-pg-upgrade:
 	$(pg_upgrade_check) --old-bindir=$(old-bindir) --new-bindir=$(new-bindir) --pgxsdir=$(pgxsdir)
 
+check-citus-upgrade:
+	$(citus_upgrade_check) \
+		--bindir=$(bindir) \
+		--pgxsdir=$(pgxsdir) \
+		--citus-pre-tar=$(citus-pre-tar) \
+		--citus-post-tar=$(citus-post-tar)
+
+check-citus-upgrade-mixed:
+	$(citus_upgrade_check) \
+		--bindir=$(bindir) \
+		--pgxsdir=$(pgxsdir) \
+		--citus-pre-tar=$(citus-pre-tar) \
+		--citus-post-tar=$(citus-post-tar) \
+		--mixed
+
 clean distclean maintainer-clean:
 	rm -f $(output_files) $(input_files)
 	rm -rf tmp_check/

src/test/regress/after_citus_upgrade_coord_schedule

@@ -0,0 +1,3 @@
+# this schedule is to be run only on coordinators
+
+test: after_citus_upgrade_coord

src/test/regress/before_citus_upgrade_coord_schedule

@@ -0,0 +1,3 @@
+# this schedule is to be run on only coordinators
+
+test: before_citus_upgrade_coord

src/test/regress/expected/after_citus_upgrade_coord.out

@@ -0,0 +1,38 @@
+SET search_path TO before_citus_upgrade_coord, public;
+SELECT * FROM t ORDER BY a;
+ a 
+---
+ 1
+ 2
+ 3
+ 4
+ 5
+(5 rows)
+
+SELECT * FROM t WHERE a = 1;
+ a 
+---
+ 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)
+
+TRUNCATE TABLE t;
+SELECT * FROM T;
+ a 
+---
+(0 rows)
+
+DROP TABLE t;
+DROP SCHEMA before_citus_upgrade_coord CASCADE;

src/test/regress/expected/before_citus_upgrade_coord.out

@@ -0,0 +1,10 @@
+CREATE SCHEMA before_citus_upgrade_coord;
+SET search_path TO before_citus_upgrade_coord, public;
+CREATE TABLE t(a int);
+SELECT create_distributed_table('t', 'a');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+INSERT INTO t SELECT * FROM generate_series(1, 5);

src/test/regress/sql/after_citus_upgrade_coord.sql

@@ -0,0 +1,17 @@
+SET search_path TO before_citus_upgrade_coord, 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;
+
+TRUNCATE TABLE t;
+
+SELECT * FROM T;
+
+DROP TABLE t;
+
+DROP SCHEMA before_citus_upgrade_coord CASCADE;

src/test/regress/sql/before_citus_upgrade_coord.sql

@@ -0,0 +1,6 @@
+CREATE SCHEMA before_citus_upgrade_coord;
+SET search_path TO before_citus_upgrade_coord, public;
+
+CREATE TABLE t(a int);
+SELECT create_distributed_table('t', 'a');
+INSERT INTO t SELECT * FROM generate_series(1, 5);

src/test/regress/upgrade/README.md

@@ -1,44 +1,83 @@
-Upgrade test is used for testing postgres version upgrade with citus installed.
+# Upgrade Tests
+
+## Postgres Upgrade Test
+
+Postgres upgrade test is used for testing postgres version upgrade with citus installed.
 Before running the script, make sure that:
+
 - You have downloaded citus.
 - You have two different postgres versions.
 - Citus is installed to both of the postgres versions. For each postgres version:
   - In citus source directory run:
-    ```
+
+  ```bash
     make clean
     ./configure PG_CONFIG=<your path to postgres pg config>
     PG_CONFIG=<your path to postgres pg config> make
     sudo PG_CONFIG=<your path to postgres pg config> make install
   ```
+
   Make sure you do this for both postgres versions, pg_config should be different for each postgres version.
+
 - Install `pipenv` and run in `citus/src/test/regress`:
-```
+
+```bash
     pipenv install
     pipenv shell
 ```
 
 - Finally run upgrade test in `citus/src/test/regress`:
-```
+
-    pipenv run make check-upgrade old-bindir=<old-bindir> new-bindir=<new-bindir>
+```bash
+    pipenv run make check-pg-upgrade old-bindir=<old-bindir> new-bindir=<new-bindir>
 ```
 
 To see full command list:
 
-```
+```bash
-    pipenv run upgrade/upgrade_test.py -help
+    pipenv run upgrade/pg_upgrade_test.py -help
 ```
 
+How the postgres upgrade test works:
 
-How the upgrade test works:
 - Temporary folder `tmp_upgrade` is created in `src/test/regress/`, if one exists it is removed first.
 - Database is initialized and citus cluster is created(1 coordinator + 2 workers) with old postgres.
-- `before_upgrade_schedule` is run with `pg_regress`. This schedule does not drop any tables or data so that we can verify upgrade.
+- `before_pg_upgrade_schedule` is run with `pg_regress`. This schedule does not drop any tables or data so that we can verify upgrade.
 - `citus_prepare_pg_upgrade` is run in coordinators and workers.
 - Old database is stopped.
 - A new database is initialized with new postgres under `tmp_upgrade`.
 - Postgres upgrade is performed.
 - New database is started in both coordinators and workers.
 - `citus_finish_pg_upgrade` is run in coordinators and workers to finalize the upgrade step.
-- `after_upgrade_schedule` is run with `pg_regress` to verify that the previously created tables, and data still exist. Router and realtime queries are used to verify this.
+- `after_pg_upgrade_schedule` is run with `pg_regress` to verify that the previously created tables, and data still exist. Router and realtime queries are used to verify this.
 
+## Citus Upgrade Test
 
+Citus upgrade test is used for testing citus version upgrades from specific version to master. The purpose of this test is to ensure that a newly made change does not result in unexpected upgrade errors.
+
+The citus upgrade test is designed to be run on a docker image for CircleCI, so we won't explain here how to run it in your local.
+
+Currently the citus upgrade test assumes that:
+
+- You have citus artifact tarballs, both for old version and master.
+
+How the citus upgrade test work:
+
+- The script takes `citus-pre-tar` and `citus-post-tar` which should contain citus artifacts.
+- It installs the given citus version from `citus-pre-tar`.
+- It creates a citus cluster(1 coordinator 2 workers).
+- It reports the initial versions.
+- It installs the checked out citus version from `citus-post-tar`.
+- It restarts the database and runs `ALTER EXTENSION citus UPGRADE`.
+- It runs `after_citus_upgrade` schedule to verify that the upgrade is successful.
+- It stops the cluster.
+
+Note that when the version of citus changes, we should update `MASTER_VERSION` with the new version of citus otherwise that will be outdated and it will fail.
+
+There is a target for citus upgrade test. We run citus upgrade tests both in normal mode and in mixed mode. In mixed mode, we dont upgrade one of the workers. `'citus.enable_version_checks' : 'false'` is used to prevent citus from giving an error for mixed mode.
+
+To see full command list:
+
+```bash
+    pipenv run upgrade/citus_upgrade_test.py -help
+```

src/test/regress/upgrade/citus_upgrade_test.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+
+"""citus_upgrade_test
+Usage:
+    citus_upgrade_test [options] --bindir=<bindir> --pgxsdir=<pgxsdir>
+
+Options:
+    --bindir=<bindir>                      The PostgreSQL executable directory(ex: '~/.pgenv/pgsql-11.3/bin')
+    --citus-pre-tar=<citus-pre-tar>        Tarball with the citus artifacts to use as the base version to upgrade from
+    --citus-post-tar=<citus-post-tar>      Tarball with the citus artifacts to use as the new version to upgrade to
+    --pgxsdir=<pgxsdir>           	       Path to the PGXS directory(ex: ~/.pgenv/src/postgresql-11.3)
+    --mixed                                Run the verification phase with one node not upgraded.
+"""
+
+import subprocess
+import atexit
+import os
+import re
+import sys
+
+import utils
+
+from docopt import docopt
+
+from config import (
+    CitusUpgradeConfig, NODE_PORTS, COORDINATOR_NAME, CITUS_VERSION_SQL, MASTER_VERSION,
+    NODE_NAMES, USER, WORKER1PORT, MASTER, HOME,
+    AFTER_CITUS_UPGRADE_COORD_SCHEDULE, BEFORE_CITUS_UPGRADE_COORD_SCHEDULE
+)
+
+import upgrade_common as common
+
+
+def main(config):
+    install_citus(config.pre_tar_path)
+    common.initialize_temp_dir(config.temp_dir)
+    common.initialize_citus_cluster(
+        config.bindir, config.datadir, config.settings)
+
+    report_initial_version(config)
+    run_test_on_coordinator(config, BEFORE_CITUS_UPGRADE_COORD_SCHEDULE)
+    remove_citus(config.pre_tar_path)
+    install_citus(config.post_tar_path)
+
+    restart_databases(config.bindir, config.datadir, config.mixed_mode)
+    run_alter_citus(config.bindir, config.mixed_mode)
+    verify_upgrade(config, config.mixed_mode)
+
+    run_test_on_coordinator(config, AFTER_CITUS_UPGRADE_COORD_SCHEDULE)
+    remove_citus(config.post_tar_path)
+
+def install_citus(tar_path):
+    with utils.cd('/'):
+        subprocess.call(['tar', 'xvf', tar_path])
+
+def report_initial_version(config):
+    for port in NODE_PORTS.values():
+        actual_citus_version = get_actual_citus_version(config.bindir, port)
+        print("port:{} citus version {}".format(port , actual_citus_version))
+
+def get_version_number(version):
+    return re.findall('\d+.\d+', version)[0]
+
+def get_actual_citus_version(pg_path, port):
+    citus_version = utils.psql(pg_path, port, CITUS_VERSION_SQL)
+    citus_version = citus_version.decode('utf-8')
+    return get_version_number(citus_version)
+
+def run_test_on_coordinator(config, schedule):
+    common.run_pg_regress(config.bindir, config.pg_srcdir,
+                   NODE_PORTS[COORDINATOR_NAME], schedule)
+
+def remove_citus(tar_path):
+    with utils.cd('/'):
+        remove_tar_files(tar_path)
+
+def remove_tar_files(tar_path):
+    ps = subprocess.Popen(('tar', 'tf', tar_path), stdout=subprocess.PIPE)
+    output = subprocess.check_output(('xargs', 'rm', '-v'), stdin=ps.stdout)
+    ps.wait()
+
+def restart_databases(pg_path, rel_data_path, mixed_mode):
+    for node_name in NODE_NAMES:
+        if mixed_mode and NODE_PORTS[node_name] == WORKER1PORT:
+            continue
+        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
+        restart_database(
+            pg_path=pg_path, abs_data_path=abs_data_path, node_name=node_name)    
+
+def restart_database(pg_path, abs_data_path, node_name):
+    command = [
+        os.path.join(pg_path, 'pg_ctl'), 'restart',
+        '--pgdata', abs_data_path,
+        '-U', USER,
+        '-o', '-p {}'.format(NODE_PORTS[node_name]),
+        '--log', os.path.join(abs_data_path, 'logfile_' + node_name)
+    ]
+    subprocess.call(command)
+
+def run_alter_citus(pg_path, mixed_mode):
+    for port in NODE_PORTS.values():
+        if mixed_mode and port == WORKER1PORT:
+            continue
+        utils.psql(pg_path, port, "ALTER EXTENSION citus UPDATE;")
+
+def verify_upgrade(config, mixed_mode):
+    for port in NODE_PORTS.values():
+        actual_citus_version = get_actual_citus_version(config.bindir, port)
+        expected_citus_version = MASTER_VERSION
+        if expected_citus_version != actual_citus_version and not (mixed_mode and port == WORKER1PORT):
+            print("port: {} citus version {} expected {}".format(port, actual_citus_version, expected_citus_version))
+            sys.exit(1)
+        else:
+            print("port:{} citus version {}".format(port , actual_citus_version))
+
+if __name__ == '__main__':
+    config = CitusUpgradeConfig(docopt(__doc__, version='citus_upgrade_test'))
+    atexit.register(common.stop_databases, config.bindir, config.datadir)
+    main(config)

src/test/regress/upgrade/config.py

@@ -1,8 +1,39 @@
-BEFORE_UPGRADE_SCHEDULE = './before_upgrade_schedule'
+from os.path import expanduser
-AFTER_UPGRADE_SCHEDULE = './after_upgrade_schedule'
 
 
-class Config():
+BEFORE_PG_UPGRADE_SCHEDULE = './before_pg_upgrade_schedule'
+AFTER_PG_UPGRADE_SCHEDULE = './after_pg_upgrade_schedule'
+
+AFTER_CITUS_UPGRADE_COORD_SCHEDULE = './after_citus_upgrade_coord_schedule'
+BEFORE_CITUS_UPGRADE_COORD_SCHEDULE = './before_citus_upgrade_coord_schedule'
+
+MASTER = 'master'
+# This should be updated when citus version changes
+MASTER_VERSION = '9.0'
+
+HOME = expanduser("~")
+
+
+CITUS_VERSION_SQL = "SELECT extversion FROM pg_extension WHERE extname = 'citus';"
+
+
+class CitusUpgradeConfig():
+    def __init__(self, arguments):
+        self.bindir = arguments['--bindir']
+        self.pre_tar_path = arguments['--citus-pre-tar']
+        self.post_tar_path = arguments['--citus-post-tar']
+        self.pg_srcdir = arguments['--pgxsdir']
+        self.temp_dir = './tmp_citus_upgrade'
+        self.datadir = self.temp_dir + '/data'
+        self.settings = {
+            'shared_preload_libraries': 'citus',
+            'citus.node_conninfo': 'sslmode=prefer',
+            'citus.enable_version_checks' : 'false'
+        }
+        self.mixed_mode = arguments['--mixed']
+
+
+class PGUpgradeConfig():
     def __init__(self, arguments):
         self.old_bindir = arguments['--old-bindir']
         self.new_bindir = arguments['--new-bindir']
@@ -23,10 +54,13 @@ COORDINATOR_NAME = 'coordinator'
 WORKER1 = 'worker1'
 WORKER2 = 'worker2'
 NODE_NAMES = [COORDINATOR_NAME, WORKER1, WORKER2]
+COORDINATOR_PORT = 57635
+WORKER1PORT = 57636
+WORKER2PORT = 57637
 
-WORKER_PORTS = [57636, 57637]
+WORKER_PORTS = [WORKER1PORT, WORKER2PORT]
 NODE_PORTS = {
-    COORDINATOR_NAME: 57635,
+    COORDINATOR_NAME: COORDINATOR_PORT,
-    WORKER1: 57636,
+    WORKER1: WORKER1PORT,
-    WORKER2: 57637,
+    WORKER2: WORKER2PORT,
 }

src/test/regress/upgrade/pg_upgrade_test.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+
+"""upgrade_test
+Usage:
+    upgrade_test --old-bindir=<old-bindir> --new-bindir=<new-bindir> --pgxsdir=<pgxsdir>
+
+Options:
+    --old-bindir=<old-bindir>              The old PostgreSQL executable directory(ex: '~/.pgenv/pgsql-10.4/bin')
+    --new-bindir=<new-bindir>              The new PostgreSQL executable directory(ex: '~/.pgenv/pgsql-11.3/bin')
+    --pgxsdir=<pgxsdir>           	       Path to the PGXS directory(ex: ~/.pgenv/src/postgresql-11.3)
+"""
+
+from config import (
+    PGUpgradeConfig, USER, NODE_PORTS,
+    NODE_NAMES, DBNAME, COORDINATOR_NAME,
+    WORKER_PORTS, AFTER_PG_UPGRADE_SCHEDULE, BEFORE_PG_UPGRADE_SCHEDULE
+)
+from docopt import docopt
+import utils
+import atexit
+import subprocess
+import sys
+import shutil
+import os
+
+import upgrade_common as common
+
+def citus_prepare_pg_upgrade(pg_path):
+    for port in NODE_PORTS.values():
+        utils.psql(pg_path, port, "SELECT citus_prepare_pg_upgrade();")
+
+
+def perform_postgres_upgrade(old_bindir, new_bindir, old_datadir, new_datadir):
+    for node_name in NODE_NAMES:
+        base_new_data_path = os.path.abspath(new_datadir)
+        base_old_data_path = os.path.abspath(old_datadir)
+        with utils.cd(base_new_data_path):
+            abs_new_data_path = os.path.join(base_new_data_path, node_name)
+            abs_old_data_path = os.path.join(base_old_data_path, node_name)
+            command = [
+                os.path.join(new_bindir, 'pg_upgrade'),
+                '--username', USER,
+                '--old-bindir', old_bindir,
+                '--new-bindir', new_bindir,
+                '--old-datadir', abs_old_data_path,
+                '--new-datadir', abs_new_data_path
+            ]
+            subprocess.call(command)
+
+
+def citus_finish_pg_upgrade(pg_path):
+    for port in NODE_PORTS.values():
+        utils.psql(pg_path, port, "SELECT citus_finish_pg_upgrade();")
+
+
+def stop_all_databases(old_bindir, new_bindir, old_datadir, new_datadir):
+    common.stop_databases(old_bindir, old_datadir)
+    common.stop_databases(new_bindir, new_datadir)
+
+
+def main(config):
+    common.initialize_temp_dir(config.temp_dir)
+    common.initialize_citus_cluster(config.old_bindir, config.old_datadir, config.settings)   
+    common.run_pg_regress(config.old_bindir, config.pg_srcdir,
+                   NODE_PORTS[COORDINATOR_NAME], BEFORE_PG_UPGRADE_SCHEDULE)
+
+    citus_prepare_pg_upgrade(config.old_bindir)
+    common.stop_databases(config.old_bindir, config.old_datadir)
+
+    common.initialize_db_for_cluster(
+        config.new_bindir, config.new_datadir, config.settings)
+    perform_postgres_upgrade(
+        config.old_bindir, config.new_bindir, config.old_datadir, config.new_datadir)
+    common.start_databases(config.new_bindir, config.new_datadir)
+    citus_finish_pg_upgrade(config.new_bindir)
+
+    common.run_pg_regress(config.new_bindir, config.pg_srcdir,
+                   NODE_PORTS[COORDINATOR_NAME], AFTER_PG_UPGRADE_SCHEDULE)
+
+
+if __name__ == '__main__':
+    config = PGUpgradeConfig(docopt(__doc__, version='upgrade_test'))
+    atexit.register(stop_all_databases, config.old_bindir,
+                    config.new_bindir, config.old_datadir, config.new_datadir)
+    main(config)

src/test/regress/upgrade/upgrade_common.py

@@ -0,0 +1,98 @@
+
+import os
+import shutil
+import sys
+import subprocess
+
+import utils
+
+from config import NODE_NAMES, NODE_PORTS, COORDINATOR_NAME, USER, WORKER_PORTS, DBNAME
+
+
+def initialize_temp_dir(temp_dir):
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+    os.mkdir(temp_dir)
+    # Give full access to TEMP_DIR so that postgres user can use it.
+    os.chmod(temp_dir, 0o777)
+
+
+def initialize_db_for_cluster(pg_path, rel_data_path, settings):
+    subprocess.call(['mkdir', rel_data_path])
+    for node_name in NODE_NAMES:
+        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
+        command = [
+            os.path.join(pg_path, 'initdb'),
+            '--pgdata', abs_data_path,
+            '--username', USER
+        ]
+        subprocess.call(command)
+        add_settings(abs_data_path, settings)
+
+
+def add_settings(abs_data_path, settings):
+    conf_path = os.path.join(abs_data_path, 'postgresql.conf')
+    with open(conf_path, 'a') as conf_file:
+        for setting_key, setting_val in settings.items():
+            setting = "{setting_key} = \'{setting_val}\'\n".format(
+                setting_key=setting_key,
+                setting_val=setting_val)
+            conf_file.write(setting)
+
+
+def start_databases(pg_path, rel_data_path):
+    for node_name in NODE_NAMES:
+        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
+        command = [
+            os.path.join(pg_path, 'pg_ctl'), 'start',
+            '--pgdata', abs_data_path,
+            '-U', USER,
+            '-o', '-p {}'.format(NODE_PORTS[node_name]),
+            '--log', os.path.join(abs_data_path, 'logfile_' + node_name)
+        ]
+        subprocess.call(command)
+
+def create_citus_extension(pg_path):
+    for port in NODE_PORTS.values():
+        utils.psql(pg_path, port, "CREATE EXTENSION citus;")
+
+def run_pg_regress(pg_path, pg_srcdir, port, schedule):
+    command = [
+        os.path.join(pg_srcdir, 'src/test/regress/pg_regress'),
+        '--port', str(port),
+        '--schedule', schedule,
+        '--bindir', pg_path,
+        '--user', USER,
+        '--dbname', DBNAME,
+        '--use-existing'
+    ]
+    exit_code = subprocess.call(command)
+    if exit_code != 0:
+        sys.exit(exit_code)
+
+
+def add_workers(pg_path):
+    for port in WORKER_PORTS:
+        command = "SELECT * from master_add_node('localhost', {port});".format(
+            port=port)
+        utils.psql(pg_path, NODE_PORTS[COORDINATOR_NAME], command)
+
+def stop_databases(pg_path, rel_data_path):
+    for node_name in NODE_NAMES:
+        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
+        command = [
+            os.path.join(pg_path, 'pg_ctl'), 'stop',
+            '--pgdata', abs_data_path,
+            '-U', USER,
+            '-o', '-p {}'.format(NODE_PORTS[node_name]),
+            '--log', os.path.join(abs_data_path, 'logfile_' + node_name)
+        ]
+        subprocess.call(command)
+
+
+def initialize_citus_cluster(old_bindir, old_datadir, settings):
+    initialize_db_for_cluster(old_bindir, old_datadir, settings)
+    start_databases(old_bindir, old_datadir)
+    create_citus_extension(old_bindir)
+    add_workers(old_bindir)
+    

src/test/regress/upgrade/upgrade_test.py

@@ -1,179 +0,0 @@
-#!/usr/bin/env python3
-
-"""upgrade_test
-Usage:
-    upgrade_test --old-bindir=<old-bindir> --new-bindir=<new-bindir> --pgxsdir=<pgxsdir>
-
-Options:
-    --old-bindir=<old-bindir>              The old PostgreSQL executable directory(ex: '~/.pgenv/pgsql-10.4/bin')
-    --new-bindir=<new-bindir>              The new PostgreSQL executable directory(ex: '~/.pgenv/pgsql-11.3/bin')
-    --pgxsdir=<pgxsdir>           	       Path to the PGXS directory(ex: ~/.pgenv/src/postgresql-11.3)
-"""
-
-import utils
-import atexit
-import subprocess
-import sys
-import shutil
-import os
-
-from docopt import docopt
-
-from config import (
-    Config, USER, NODE_PORTS,
-    NODE_NAMES, DBNAME, COORDINATOR_NAME,
-    WORKER_PORTS, AFTER_UPGRADE_SCHEDULE, BEFORE_UPGRADE_SCHEDULE
-)
-
-
-def initialize_temp_dir(temp_dir):
-    if os.path.exists(temp_dir):
-        shutil.rmtree(temp_dir)
-    os.mkdir(temp_dir)
-    # Give full access to TEMP_DIR so that postgres user can use it.
-    os.chmod(temp_dir, 0o777)
-
-
-def initialize_db_for_cluster(pg_path, rel_data_path, settings):
-    subprocess.call(['mkdir', rel_data_path])
-    for node_name in NODE_NAMES:
-        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
-        command = [
-            os.path.join(pg_path, 'initdb'),
-            '--pgdata', abs_data_path,
-            '--username', USER
-        ]
-        subprocess.call(command)
-        add_settings(abs_data_path, settings)
-
-
-def add_settings(abs_data_path, settings):
-    conf_path = os.path.join(abs_data_path, 'postgresql.conf')
-    with open(conf_path, 'a') as conf_file:
-        for setting_key, setting_val in settings.items():
-            setting = "{setting_key} = \'{setting_val}\'\n".format(
-                setting_key=setting_key,
-                setting_val=setting_val)
-            conf_file.write(setting)
-
-
-def start_databases(pg_path, rel_data_path):
-    for node_name in NODE_NAMES:
-        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
-        command = [
-            os.path.join(pg_path, 'pg_ctl'), 'start',
-            '--pgdata', abs_data_path,
-            '-U', USER,
-            '-o', '-p {}'.format(NODE_PORTS[node_name]),
-            '--log', os.path.join(abs_data_path, 'logfile_' + node_name)
-        ]
-        subprocess.call(command)
-
-
-def create_citus_extension(pg_path):
-    for port in NODE_PORTS.values():
-        utils.psql(pg_path, port, "CREATE EXTENSION citus;")
-
-
-def add_workers(pg_path):
-    for port in WORKER_PORTS:
-        command = "SELECT * from master_add_node('localhost', {port});".format(
-            port=port)
-        utils.psql(pg_path, NODE_PORTS[COORDINATOR_NAME], command)
-
-
-def run_pg_regress(pg_path, PG_SRCDIR, port, schedule):
-    command = [
-        os.path.join(PG_SRCDIR, 'src/test/regress/pg_regress'),
-        '--port', str(port),
-        '--schedule', schedule,
-        '--bindir', pg_path,
-        '--user', USER,
-        '--dbname', DBNAME,
-        '--use-existing'
-    ]
-    exit_code = subprocess.call(command)
-    if exit_code != 0:
-        sys.exit(exit_code)
-
-
-def citus_prepare_pg_upgrade(pg_path):
-    for port in NODE_PORTS.values():
-        utils.psql(pg_path, port, "SELECT citus_prepare_pg_upgrade();")
-
-
-def stop_databases(pg_path, rel_data_path):
-    for node_name in NODE_NAMES:
-        abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
-        command = [
-            os.path.join(pg_path, 'pg_ctl'), 'stop',
-            '--pgdata', abs_data_path,
-            '-U', USER,
-            '-o', '-p {}'.format(NODE_PORTS[node_name]),
-            '--log', os.path.join(abs_data_path, 'logfile_' + node_name)
-        ]
-        subprocess.call(command)
-
-
-def perform_postgres_upgrade(old_bindir, new_bindir, old_datadir, new_datadir):
-    for node_name in NODE_NAMES:
-        base_new_data_path = os.path.abspath(new_datadir)
-        base_old_data_path = os.path.abspath(old_datadir)
-        with utils.cd(base_new_data_path):
-            abs_new_data_path = os.path.join(base_new_data_path, node_name)
-            abs_old_data_path = os.path.join(base_old_data_path, node_name)
-            command = [
-                os.path.join(new_bindir, 'pg_upgrade'),
-                '--username', USER,
-                '--old-bindir', old_bindir,
-                '--new-bindir', new_bindir,
-                '--old-datadir', abs_old_data_path,
-                '--new-datadir', abs_new_data_path
-            ]
-            subprocess.call(command)
-
-
-def citus_finish_pg_upgrade(pg_path):
-    for port in NODE_PORTS.values():
-        utils.psql(pg_path, port, "SELECT citus_finish_pg_upgrade();")
-
-
-def initialize_citus_cluster(old_bindir, old_datadir, settings):
-    initialize_db_for_cluster(old_bindir, old_datadir, settings)
-    start_databases(old_bindir, old_datadir)
-    create_citus_extension(old_bindir)
-    add_workers(old_bindir)
-
-
-def stop_all_databases(old_bindir, new_bindir, old_datadir, new_datadir):
-    stop_databases(old_bindir, old_datadir)
-    stop_databases(new_bindir, new_datadir)
-
-
-def main(config):
-    initialize_temp_dir(config.temp_dir)
-    initialize_citus_cluster(
-        config.old_bindir, config.old_datadir, config.settings)
-
-    run_pg_regress(config.old_bindir, config.pg_srcdir,
-                   NODE_PORTS[COORDINATOR_NAME], BEFORE_UPGRADE_SCHEDULE)
-
-    citus_prepare_pg_upgrade(config.old_bindir)
-    stop_databases(config.old_bindir, config.old_datadir)
-
-    initialize_db_for_cluster(
-        config.new_bindir, config.new_datadir, config.settings)
-    perform_postgres_upgrade(
-        config.old_bindir, config.new_bindir, config.old_datadir, config.new_datadir)
-    start_databases(config.new_bindir, config.new_datadir)
-    citus_finish_pg_upgrade(config.new_bindir)
-
-    run_pg_regress(config.new_bindir, config.pg_srcdir,
-                   NODE_PORTS[COORDINATOR_NAME], AFTER_UPGRADE_SCHEDULE)
-
-
-if __name__ == '__main__':
-    config = Config(docopt(__doc__, version='upgrade_test'))
-    atexit.register(stop_all_databases, config.old_bindir,
-                    config.new_bindir, config.old_datadir, config.new_datadir)
-    main(config)

src/test/regress/upgrade/utils.py

@@ -4,8 +4,7 @@ from config import USER
 
 
 def psql(pg_path, port, command):
-
+    return subprocess.check_output([
-    return subprocess.call([
         os.path.join(pg_path, 'psql'),
         '-U', USER,
         '-p', str(port),