Make python-regress based tests runnable with run_test.py (#6814)

For some tests such as upgrade tests and arbitrary config tests we set
up the citus cluster using Python. This setup is slightly different from
the perl based setup script (`multi_regress.pl`). Most importantly it
uses replication factor 1 by default.

This changes our run_test.py script to be able to run a schedule using
python instead of `multi_regress.pl`, for the tests that require it.

For now arbitrary config tests are still not runnable with
`run_test.py`, but this brings us one step closer to being able to do
that.

Fixes #6804

src/test/regress/citus_tests/common.py

@@ -360,6 +360,11 @@ def run(command, *args, check=True, shell=True, silent=False, **kwargs):
     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/run_test.py

@@ -13,7 +13,7 @@ from typing import Optional
 
 import common
 
-import config
+from config import ARBITRARY_SCHEDULE_NAMES, MASTER_VERSION, CitusDefaultClusterConfig
 
 
 # Returns true if given test_schedule_line is of the form:
@@ -46,6 +46,30 @@ def run_python_test(test_file_name, repeat):
     )
 
 
+def run_schedule_with_python(schedule):
+    bindir = common.capture("pg_config --bindir").rstrip()
+    pgxs_path = pathlib.Path(common.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 = {
+        "--pgxsdir": str(pgxs_path.parent.parent.parent),
+        "--bindir": bindir,
+    }
+
+    config = CitusDefaultClusterConfig(args)
+    common.initialize_temp_dir(config.temp_dir)
+    common.initialize_citus_cluster(
+        config.bindir, config.datadir, config.settings, config
+    )
+    common.run_pg_regress(
+        config.bindir, config.pg_srcdir, config.coordinator_port(), schedule
+    )
+
+
 if __name__ == "__main__":
     args = argparse.ArgumentParser()
     args.add_argument(
@@ -210,7 +234,19 @@ if __name__ == "__main__":
         if "operations" in test_schedule:
             return "minimal_schedule"
 
-        if test_schedule in config.ARBITRARY_SCHEDULE_NAMES:
+        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_schedule"
+
+        if test_schedule in ARBITRARY_SCHEDULE_NAMES:
             print(
                 f"WARNING: Arbitrary config schedule ({test_schedule}) is not supported."
             )
@@ -239,6 +275,9 @@ if __name__ == "__main__":
     else:
         dependencies = TestDeps(default_base_schedule(test_schedule))
 
+    if "before_" in test_schedule:
+        dependencies.repeatable = False
+
     # copy base schedule to a temp file and append test_schedule_line
     # to be able to run tests in parallel (if test_schedule_line is a parallel group.)
     tmp_schedule_path = os.path.join(
@@ -262,6 +301,14 @@ if __name__ == "__main__":
         for _ in range(repetition_cnt):
             myfile.write(test_schedule_line)
 
+    if "upgrade" in test_schedule_line:
+        try:
+            run_schedule_with_python(pathlib.Path(tmp_schedule_path).stem)
+        finally:
+            # remove temp schedule file
+            os.remove(tmp_schedule_path)
+        sys.exit(0)
+
     # find suitable make recipe
     if dependencies.schedule == "base_isolation_schedule":
         make_recipe = "check-isolation-custom-schedule"