Support running Citus upgrade tests with run_test.py (#6832)

Citus upgrade tests require some additional logic to run, because we
have a before and after schedule and we need to swap the Citus
version in-between. This adds that logic to `run_test.py`.

In passing this makes running upgrade tests locally multiple times
faster by caching tarballs.

.flake8

@@ -4,3 +4,4 @@ extend-ignore = E203
 # black will truncate to 88 characters usually, but long string literals it
 # might keep. That's fine in most cases unless it gets really excessive.
 max-line-length = 150
+exclude = .git,__pycache__,vendor,tmp_*

src/test/regress/Makefile

@@ -295,7 +295,7 @@ check-citus-upgrade-mixed-local: all clean-upgrade-artifacts
 		--mixed
 
 clean-upgrade-artifacts:
-	rm -rf $(citus_abs_srcdir)/tmp_citus_tarballs/ $(citus_abs_srcdir)/tmp_citus_upgrade/ /tmp/citus_copy/
+	rm -rf $(citus_abs_srcdir)/tmp_citus_upgrade/ /tmp/citus_copy/
 
 clean distclean maintainer-clean:
 	rm -rf input/ output/

src/test/regress/citus_tests/common.py

@@ -48,6 +48,54 @@ TIMEOUT_DEFAULT = timedelta(seconds=int(os.getenv("PG_TEST_TIMEOUT_DEFAULT", "10
 FORCE_PORTS = os.getenv("PG_FORCE_PORTS", "NO").lower() not in ("no", "0", "n", "")
 
 REGRESS_DIR = pathlib.Path(os.path.realpath(__file__)).parent.parent
+REPO_ROOT = REGRESS_DIR.parent.parent.parent
+CI = os.environ.get("CI") == "true"
+
+
+def eprint(*args, **kwargs):
+    """eprint prints to stderr"""
+
+    print(*args, file=sys.stderr, **kwargs)
+
+
+def run(command, *args, check=True, shell=True, silent=False, **kwargs):
+    """run runs the given command and prints it to stderr"""
+
+    if not silent:
+        eprint(f"+ {command} ")
+    if silent:
+        kwargs.setdefault("stdout", subprocess.DEVNULL)
+    return subprocess.run(command, *args, check=check, shell=shell, **kwargs)
+
+
+def capture(command, *args, **kwargs):
+    """runs the given command and returns its output as a string"""
+    return run(command, *args, stdout=subprocess.PIPE, text=True, **kwargs).stdout
+
+
+PG_CONFIG = os.environ.get("PG_CONFIG", "pg_config")
+PG_BINDIR = capture([PG_CONFIG, "--bindir"], shell=False).rstrip()
+os.environ["PATH"] = PG_BINDIR + os.pathsep + os.environ["PATH"]
+
+
+def get_pg_major_version():
+    full_version_string = run(
+        "initdb --version", stdout=subprocess.PIPE, encoding="utf-8", silent=True
+    ).stdout
+    major_version_string = re.search("[0-9]+", full_version_string)
+    assert major_version_string is not None
+    return int(major_version_string.group(0))
+
+
+PG_MAJOR_VERSION = get_pg_major_version()
+
+OLDEST_SUPPORTED_CITUS_VERSION_MATRIX = {
+    13: "9.5.0",
+    14: "10.2.0",
+    15: "11.1.5",
+}
+
+OLDEST_SUPPORTED_CITUS_VERSION = OLDEST_SUPPORTED_CITUS_VERSION_MATRIX[PG_MAJOR_VERSION]
 
 
 def initialize_temp_dir(temp_dir):
@@ -357,27 +405,6 @@ def initialize_citus_cluster(bindir, datadir, settings, config):
     config.setup_steps()
 
 
-def eprint(*args, **kwargs):
-    """eprint prints to stderr"""
-
-    print(*args, file=sys.stderr, **kwargs)
-
-
-def run(command, *args, check=True, shell=True, silent=False, **kwargs):
-    """run runs the given command and prints it to stderr"""
-
-    if not silent:
-        eprint(f"+ {command} ")
-    if silent:
-        kwargs.setdefault("stdout", subprocess.DEVNULL)
-    return subprocess.run(command, *args, check=check, shell=shell, **kwargs)
-
-
-def capture(command, *args, **kwargs):
-    """runs the given command and returns its output as a string"""
-    return run(command, *args, stdout=subprocess.PIPE, text=True, **kwargs).stdout
-
-
 def sudo(command, *args, shell=True, **kwargs):
     """
     A version of run that prefixes the command with sudo when the process is

src/test/regress/citus_tests/config.py

@@ -179,10 +179,10 @@ class CitusDefaultClusterConfig(CitusBaseClusterConfig):
 
 
 class CitusUpgradeConfig(CitusBaseClusterConfig):
-    def __init__(self, arguments):
+    def __init__(self, arguments, pre_tar, post_tar):
         super().__init__(arguments)
-        self.pre_tar_path = arguments["--citus-pre-tar"]
-        self.post_tar_path = arguments["--citus-post-tar"]
+        self.pre_tar_path = pre_tar
+        self.post_tar_path = post_tar
         self.temp_dir = "./tmp_citus_upgrade"
         self.new_settings = {"citus.enable_version_checks": "false"}
         self.user = SUPER_USER_NAME

src/test/regress/citus_tests/run_test.py

@@ -13,9 +13,10 @@ from contextlib import contextmanager
 from typing import Optional
 
 import common
-from common import REGRESS_DIR, capture, run
+from common import OLDEST_SUPPORTED_CITUS_VERSION, PG_BINDIR, REGRESS_DIR, capture, run
+from upgrade import generate_citus_tarball, run_citus_upgrade_tests
 
-from config import ARBITRARY_SCHEDULE_NAMES, MASTER_VERSION, CitusBaseClusterConfig
+from config import ARBITRARY_SCHEDULE_NAMES, CitusBaseClusterConfig, CitusUpgradeConfig
 
 
 def main():
@@ -75,11 +76,19 @@ class TestDeps:
     schedule: Optional[str]
     direct_extra_tests: list[str]
 
-    def __init__(self, schedule, extra_tests=None, repeatable=True, worker_count=2):
+    def __init__(
+        self,
+        schedule,
+        extra_tests=None,
+        repeatable=True,
+        worker_count=2,
+        citus_upgrade_infra=False,
+    ):
         self.schedule = schedule
         self.direct_extra_tests = extra_tests or []
         self.repeatable = repeatable
         self.worker_count = worker_count
+        self.citus_upgrade_infra = citus_upgrade_infra
 
     def extra_tests(self):
         all_deps = OrderedDict()
@@ -176,26 +185,29 @@ def run_regress_test(test_name, args):
 
     with tmp_schedule(test_name, dependencies, schedule_line, args) as schedule:
         if "upgrade" in original_schedule:
-            run_schedule_with_python(schedule)
+            run_schedule_with_python(test_name, schedule, dependencies)
         else:
             run_schedule_with_multiregress(test_name, schedule, dependencies, args)
 
 
-def run_schedule_with_python(schedule):
-    bindir = capture("pg_config --bindir").rstrip()
+def run_schedule_with_python(test_name, schedule, dependencies):
     pgxs_path = pathlib.Path(capture("pg_config --pgxs").rstrip())
 
     os.chdir(REGRESS_DIR)
     os.environ["PATH"] = str(REGRESS_DIR / "bin") + os.pathsep + os.environ["PATH"]
     os.environ["PG_REGRESS_DIFF_OPTS"] = "-dU10 -w"
-    os.environ["CITUS_OLD_VERSION"] = f"v{MASTER_VERSION}.0"
 
-    args = {
+    fake_config_args = {
         "--pgxsdir": str(pgxs_path.parent.parent.parent),
-        "--bindir": bindir,
+        "--bindir": PG_BINDIR,
+        "--mixed": False,
     }
 
-    config = CitusBaseClusterConfig(args)
+    if dependencies.citus_upgrade_infra:
+        run_single_citus_upgrade_test(test_name, schedule, fake_config_args)
+        return
+
+    config = CitusBaseClusterConfig(fake_config_args)
     common.initialize_temp_dir(config.temp_dir)
     common.initialize_citus_cluster(
         config.bindir, config.datadir, config.settings, config
@@ -205,6 +217,29 @@ def run_schedule_with_python(schedule):
     )
 
 
+def run_single_citus_upgrade_test(test_name, schedule, fake_config_args):
+    os.environ["CITUS_OLD_VERSION"] = f"v{OLDEST_SUPPORTED_CITUS_VERSION}"
+    citus_tarball_path = generate_citus_tarball(f"v{OLDEST_SUPPORTED_CITUS_VERSION}")
+    config = CitusUpgradeConfig(fake_config_args, citus_tarball_path, None)
+
+    # Before tests are a simple case, because no actual upgrade is needed
+    if "_before" in test_name:
+        run_citus_upgrade_tests(config, schedule, None)
+        return
+
+    before_schedule_name = f"{schedule}_before"
+    before_schedule_path = REGRESS_DIR / before_schedule_name
+
+    before_test_name = test_name.replace("_after", "_before")
+    with open(before_schedule_path, "w") as before_schedule_file:
+        before_schedule_file.write(f"test: {before_test_name}\n")
+
+    try:
+        run_citus_upgrade_tests(config, before_schedule_name, schedule)
+    finally:
+        os.remove(before_schedule_path)
+
+
 def run_schedule_with_multiregress(test_name, schedule, dependencies, args):
     worker_count = needed_worker_count(test_name, dependencies)
 
@@ -249,15 +284,6 @@ def default_base_schedule(test_schedule, args):
     if "operations" in test_schedule:
         return "minimal_schedule"
 
-    if "after_citus_upgrade" in test_schedule:
-        print(
-            f"WARNING: After citus upgrade schedule ({test_schedule}) is not supported."
-        )
-        sys.exit(0)
-
-    if "citus_upgrade" in test_schedule:
-        return None
-
     if "pg_upgrade" in test_schedule:
         return "minimal_pg_upgrade_schedule"
 
@@ -316,6 +342,9 @@ def test_dependencies(test_name, test_schedule, schedule_line, args):
     if test_name in DEPS:
         return DEPS[test_name]
 
+    if "citus_upgrade" in test_schedule:
+        return TestDeps(None, citus_upgrade_infra=True)
+
     if schedule_line_is_upgrade_after(schedule_line):
         # upgrade_xxx_after tests always depend on upgrade_xxx_before
         test_names = schedule_line.split()[1:]

src/test/regress/citus_tests/upgrade/__init__.py

@@ -0,0 +1 @@
+from .citus_upgrade_test import generate_citus_tarball, run_citus_upgrade_tests  # noqa

src/test/regress/citus_tests/upgrade/citus_upgrade_test.py

@@ -13,7 +13,7 @@ Options:
     --mixed                                 Run the verification phase with one node not upgraded.
 """
 
-import atexit
+import multiprocessing
 import os
 import re
 import subprocess
@@ -27,6 +27,7 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 # ignore E402 because these imports require addition to path
 import common  # noqa: E402
 import utils  # noqa: E402
+from common import CI, PG_MAJOR_VERSION, REPO_ROOT, run  # noqa: E402
 from utils import USER  # noqa: E402
 
 from config import (  # noqa: E402
@@ -41,6 +42,12 @@ from config import (  # noqa: E402
 
 
 def main(config):
+    before_upgrade_schedule = get_before_upgrade_schedule(config.mixed_mode)
+    after_upgrade_schedule = get_after_upgrade_schedule(config.mixed_mode)
+    run_citus_upgrade_tests(config, before_upgrade_schedule, after_upgrade_schedule)
+
+
+def run_citus_upgrade_tests(config, before_upgrade_schedule, after_upgrade_schedule):
     install_citus(config.pre_tar_path)
     common.initialize_temp_dir(config.temp_dir)
     common.initialize_citus_cluster(
@@ -48,23 +55,28 @@ def main(config):
     )
 
     report_initial_version(config)
-    before_upgrade_schedule = get_before_upgrade_schedule(config.mixed_mode)
     run_test_on_coordinator(config, before_upgrade_schedule)
     remove_citus(config.pre_tar_path)
+    if after_upgrade_schedule is None:
+        return
+
     install_citus(config.post_tar_path)
 
     restart_databases(config.bindir, config.datadir, config.mixed_mode, config)
     run_alter_citus(config.bindir, config.mixed_mode, config)
     verify_upgrade(config, config.mixed_mode, config.node_name_to_ports.values())
 
-    after_upgrade_schedule = get_after_upgrade_schedule(config.mixed_mode)
     run_test_on_coordinator(config, after_upgrade_schedule)
     remove_citus(config.post_tar_path)
 
 
 def install_citus(tar_path):
-    with utils.cd("/"):
-        subprocess.run(["tar", "xvf", tar_path], check=True)
+    if tar_path:
+        with utils.cd("/"):
+            run(["tar", "xvf", tar_path], shell=False)
+    else:
+        with utils.cd(REPO_ROOT):
+            run(f"make -j{multiprocessing.cpu_count()} -s install")
 
 
 def report_initial_version(config):
@@ -90,8 +102,9 @@ def run_test_on_coordinator(config, schedule):
 
 
 def remove_citus(tar_path):
-    with utils.cd("/"):
-        remove_tar_files(tar_path)
+    if tar_path:
+        with utils.cd("/"):
+            remove_tar_files(tar_path)
 
 
 def remove_tar_files(tar_path):
@@ -171,43 +184,29 @@ def get_after_upgrade_schedule(mixed_mode):
         return AFTER_CITUS_UPGRADE_COORD_SCHEDULE
 
 
-# IsRunningOnLocalMachine returns true if the upgrade test is run on
-# local machine, in which case the old citus version will be installed
-# and it will be upgraded to the current code.
-def IsRunningOnLocalMachine(arguments):
-    return arguments["--citus-old-version"]
-
-
-def generate_citus_tarballs(citus_version):
+def generate_citus_tarball(citus_version):
     tmp_dir = "tmp_citus_tarballs"
     citus_old_tarpath = os.path.abspath(
-        os.path.join(tmp_dir, "install-citus{}.tar".format(citus_version))
-    )
-    citus_new_tarpath = os.path.abspath(
-        os.path.join(tmp_dir, "install-citusmaster.tar")
+        os.path.join(tmp_dir, f"install-pg{PG_MAJOR_VERSION}-citus{citus_version}.tar")
     )
 
     common.initialize_temp_dir_if_not_exists(tmp_dir)
     dirpath = os.path.dirname(os.path.realpath(__file__))
     local_script_path = os.path.join(dirpath, "generate_citus_tarballs.sh")
     with utils.cd(tmp_dir):
-        subprocess.check_call([local_script_path, citus_version])
+        subprocess.check_call([local_script_path, str(PG_MAJOR_VERSION), citus_version])
 
-    return [citus_old_tarpath, citus_new_tarpath]
+    return citus_old_tarpath
 
 
 if __name__ == "__main__":
     args = docopt(__doc__, version="citus_upgrade_test")
-    if IsRunningOnLocalMachine(args):
-        citus_tarball_paths = generate_citus_tarballs(args["--citus-old-version"])
-        args["--citus-pre-tar"] = citus_tarball_paths[0]
-        args["--citus-post-tar"] = citus_tarball_paths[1]
-    config = CitusUpgradeConfig(args)
-    atexit.register(
-        common.stop_databases,
-        config.bindir,
-        config.datadir,
-        config.node_name_to_ports,
-        config.name,
-    )
+    if not CI:
+        citus_tarball_path = generate_citus_tarball(args["--citus-old-version"])
+        config = CitusUpgradeConfig(args, citus_tarball_path, None)
+    else:
+        config = CitusUpgradeConfig(
+            args, args["--citus-pre-tar"], args["--citus-post-tar"]
+        )
+
     main(config)

src/test/regress/citus_tests/upgrade/generate_citus_tarballs.sh

@@ -2,7 +2,8 @@
 
 set -euxo pipefail
 
-citus_old_version=$1
+pg_version=$1
+citus_old_version=$2
 
 base="$(pwd)"
 
@@ -17,38 +18,22 @@ install_citus_and_tar() {
 
     cd "${installdir}" && find . -type f -print >"${builddir}/files.lst"
 
-    tar cvf "${basedir}/install-citus${citus_version}.tar" $(cat "${builddir}"/files.lst)
-    mv "${basedir}/install-citus${citus_version}.tar" "${base}/install-citus${citus_version}.tar"
+    tar cvf "${basedir}/install-pg${pg_version}-citus${citus_version}.tar" $(cat "${builddir}"/files.lst)
+    mv "${basedir}/install-pg${pg_version}-citus${citus_version}.tar" "${base}/install-pg${pg_version}-citus${citus_version}.tar"
 
     cd "${builddir}" && rm -rf install files.lst && make clean
 }
 
-build_current() {
-    citus_version="$1"
-    basedir="${base}/${citus_version}"
-
-    mkdir -p "${basedir}"
-    citus_repo=$(git rev-parse --show-toplevel)
-
-    cd "$citus_repo" && cp -R . /tmp/citus_copy
-    # https://stackoverflow.com/questions/957928/is-there-a-way-to-get-the-git-root-directory-in-one-command
-    mv /tmp/citus_copy "${basedir}/citus_${citus_version}"
-    builddir="${basedir}/build"
-    cd "${basedir}"
-
-    citus_dir=${basedir}/citus_$citus_version
-
-    make -C "${citus_dir}" clean
-    cd "${citus_dir}"
-    ./configure --without-libcurl
-
-    install_citus_and_tar
-}
-
 build_ext() {
     citus_version="$1"
+    # If tarball already exsists we're good
+    if [ -f "${base}/install-pg${pg_version}-citus${citus_version}.tar" ]; then
+        return
+    fi
+
     basedir="${base}/${citus_version}"
 
+    rm -rf "${basedir}"
     mkdir -p "${basedir}"
     cd "${basedir}"
     citus_dir=${basedir}/citus_$citus_version
@@ -58,5 +43,4 @@ build_ext() {
     install_citus_and_tar
 }
 
-build_current "master"
 build_ext "${citus_old_version}"