mirror of
https://github.com/apple/foundationdb.git
synced 2025-06-02 03:12:12 +08:00
78 lines
2.5 KiB
Python
Executable File
78 lines
2.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from local_cluster import random_secret_string
|
|
from argparse import ArgumentParser, RawDescriptionHelpFormatter
|
|
from pathlib import Path
|
|
|
|
|
|
class ClusterFileGenerator:
|
|
def __init__(self, output_dir: str):
|
|
self.output_dir = Path(output_dir).resolve()
|
|
assert self.output_dir.exists(), "{} does not exist".format(output_dir)
|
|
assert self.output_dir.is_dir(), "{} is not a directory".format(output_dir)
|
|
self.tmp_dir = self.output_dir.joinpath("tmp", random_secret_string(16))
|
|
self.tmp_dir.mkdir(parents=True)
|
|
self.cluster_file_path = self.tmp_dir.joinpath("fdb.cluster")
|
|
|
|
def __enter__(self):
|
|
with open(self.cluster_file_path, "x") as f:
|
|
f.write("foo:bar@1.1.1.1:5678\n")
|
|
|
|
return self
|
|
|
|
def __exit__(self, xc_type, exc_value, traceback):
|
|
shutil.rmtree(self.tmp_dir)
|
|
|
|
def close(self):
|
|
shutil.rmtree(self.tmp_dir)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser(
|
|
formatter_class=RawDescriptionHelpFormatter,
|
|
description="""
|
|
This script generates a cluster file that can be used to run a test and that will
|
|
be cleaned up when the test is over. The cluster file will not correspond to a real
|
|
cluster.
|
|
|
|
Before the command is executed, the following arguments will be preprocessed:
|
|
- All occurrences of @CLUSTER_FILE@ will be replaced with the path to the generated cluster file.
|
|
|
|
The environment variable FDB_CLUSTER_FILE is set to the generated cluster file for the command if
|
|
it is not set already.
|
|
""",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
"-o",
|
|
metavar="OUTPUT_DIRECTORY",
|
|
help="Directory where output files are written",
|
|
required=True,
|
|
)
|
|
parser.add_argument("cmd", metavar="COMMAND", nargs="+", help="The command to run")
|
|
args = parser.parse_args()
|
|
errcode = 1
|
|
|
|
with ClusterFileGenerator(args.output_dir) as generator:
|
|
print("cluster-file: {}".format(generator.cluster_file_path))
|
|
cmd_args = []
|
|
for cmd in args.cmd:
|
|
if cmd == "@CLUSTER_FILE@":
|
|
cmd_args.append(str(generator.cluster_file_path))
|
|
else:
|
|
cmd_args.append(cmd)
|
|
|
|
env = dict(**os.environ)
|
|
env["FDB_CLUSTER_FILE"] = env.get(
|
|
"FDB_CLUSTER_FILE", generator.cluster_file_path
|
|
)
|
|
errcode = subprocess.run(
|
|
cmd_args, stdout=sys.stdout, stderr=sys.stderr, env=env
|
|
).returncode
|
|
|
|
sys.exit(errcode)
|