diff --git a/scripts/export_import_between_pageservers.py b/scripts/export_import_between_pageservers.py index 152ce40cea..1734038661 100755 --- a/scripts/export_import_between_pageservers.py +++ b/scripts/export_import_between_pageservers.py @@ -370,7 +370,7 @@ def pack_base(log_dir, restored_dir, output_tar): shutil.move(tmp_tar_path, output_tar) -def reconstruct_paths(log_dir, pg_bin, base_tar): +def reconstruct_paths(log_dir, pg_bin, base_tar, port: int): """Reconstruct what relation files should exist in the datadir by querying postgres.""" with tempfile.TemporaryDirectory() as restored_dir: # Unpack the base tar @@ -378,8 +378,7 @@ def reconstruct_paths(log_dir, pg_bin, base_tar): # Start a vanilla postgres from the given datadir and query it to find # what relfiles should exist, but possibly don't. - port = "55439" # Probably free - with VanillaPostgres(restored_dir, pg_bin, port, init=False) as vanilla_pg: + with VanillaPostgres(Path(restored_dir), pg_bin, port, init=False) as vanilla_pg: vanilla_pg.configure([f"port={port}"]) vanilla_pg.start(log_path=os.path.join(log_dir, "tmp_pg.log")) @@ -443,8 +442,8 @@ def touch_missing_rels(log_dir, corrupt_tar, output_tar, paths): # a vanilla postgres from the exported datadir, and query it # to see what empty relations are missing, and then create # those empty files before importing. -def add_missing_rels(base_tar, output_tar, log_dir, pg_bin): - reconstructed_paths = set(reconstruct_paths(log_dir, pg_bin, base_tar)) +def add_missing_rels(base_tar, output_tar, log_dir, pg_bin, tmp_pg_port: int): + reconstructed_paths = set(reconstruct_paths(log_dir, pg_bin, base_tar, tmp_pg_port)) touch_missing_rels(log_dir, base_tar, output_tar, reconstructed_paths) @@ -535,7 +534,7 @@ def export_timeline( # Add missing rels pg_bin = PgBin(args.work_dir, args.pg_distrib_dir, pg_version) - add_missing_rels(incomplete_filename, tar_filename, args.work_dir, pg_bin) + add_missing_rels(incomplete_filename, tar_filename, args.work_dir, pg_bin, args.tmp_pg_port) # Log more info file_size = os.path.getsize(tar_filename) @@ -633,6 +632,13 @@ def main(args: argparse.Namespace): raise AssertionError(f"Sizes don't match old: {old_size} new: {new_size}") +def non_zero_tcp_port(arg: Any): + port = int(arg) + if port < 1 or port > 65535: + raise argparse.ArgumentTypeError(f"invalid tcp port: {arg}") + return port + + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( @@ -728,5 +734,13 @@ if __name__ == "__main__": default=False, help="directory where temporary tar files are stored", ) + parser.add_argument( + "--tmp-pg-port", + dest="tmp_pg_port", + required=False, + default=55439, + type=non_zero_tcp_port, + help="localhost port to use for temporary postgres instance", + ) args = parser.parse_args() main(args) diff --git a/scripts/reformat b/scripts/reformat index 67140a705a..5346c78ead 100755 --- a/scripts/reformat +++ b/scripts/reformat @@ -6,6 +6,6 @@ set -euox pipefail echo 'Reformatting Rust code' cargo fmt echo 'Reformatting Python code' -poetry run isort test_runner -poetry run flake8 test_runner -poetry run black test_runner +poetry run isort test_runner scripts +poetry run flake8 test_runner scripts +poetry run black test_runner scripts diff --git a/test_runner/regress/test_tenant_relocation.py b/test_runner/regress/test_tenant_relocation.py index 2c11812a7b..fa00a4da82 100644 --- a/test_runner/regress/test_tenant_relocation.py +++ b/test_runner/regress/test_tenant_relocation.py @@ -392,6 +392,8 @@ def test_tenant_relocation( pg_distrib_dir, "--work-dir", os.path.join(test_output_dir), + "--tmp-pg-port", + str(port_distributor.get_port()), ] subprocess_capture(test_output_dir, cmd, check=True) elif method == "minor":