Content-Length: 5166 | pFad | http://github.com/postgrespro/testgres/pull/131.patch

thub.com From 800e648fc30710dc779dcc056441159bf8f8d810 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Tue, 25 Jun 2024 17:24:24 +0700 Subject: [PATCH] Add an SSH port parameter --- testgres/operations/os_ops.py | 3 ++- testgres/operations/remote_ops.py | 26 +++++++++++++++----------- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py index dd6613cf..236a08c6 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -10,8 +10,9 @@ class ConnectionParams: - def __init__(self, host='127.0.0.1', ssh_key=None, username=None): + def __init__(self, host='127.0.0.1', port=None, ssh_key=None, username=None): self.host = host + self.port = port self.ssh_key = ssh_key self.username = username diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 01251e1c..697b4258 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -44,11 +44,13 @@ def __init__(self, conn_params: ConnectionParams): super().__init__(conn_params.username) self.conn_params = conn_params self.host = conn_params.host + self.port = conn_params.port self.ssh_key = conn_params.ssh_key + self.ssh_args = [] if self.ssh_key: - self.ssh_cmd = ["-i", self.ssh_key] - else: - self.ssh_cmd = [] + self.ssh_args += ["-i", self.ssh_key] + if self.port: + self.ssh_args += ["-p", self.port] self.remote = True self.username = conn_params.username or self.get_user() self.add_known_host(self.host) @@ -95,9 +97,9 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, """ ssh_cmd = [] if isinstance(cmd, str): - ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + [cmd] + ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_args + [cmd] elif isinstance(cmd, list): - ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + cmd + ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_args + cmd process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if get_process: return process @@ -246,9 +248,9 @@ def mkdtemp(self, prefix=None): - prefix (str): The prefix of the temporary directory name. """ if prefix: - command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", f"mktemp -d {prefix}XXXXX"] + command = ["ssh"] + self.ssh_args + [f"{self.username}@{self.host}", f"mktemp -d {prefix}XXXXX"] else: - command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", "mktemp -d"] + command = ["ssh"] + self.ssh_args + [f"{self.username}@{self.host}", "mktemp -d"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -291,8 +293,10 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal mode = "r+b" if binary else "r+" with tempfile.NamedTemporaryFile(mode=mode, delete=False) as tmp_file: + # For scp the port is specified by a "-P" option + scp_args = ['-P' if x == '-p' else x for x in self.ssh_args] if not truncate: - scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.username}@{self.host}:{filename}", tmp_file.name] + scp_cmd = ['scp'] + scp_args + [f"{self.username}@{self.host}:{filename}", tmp_file.name] subprocess.run(scp_cmd, check=False) # The file might not exist yet tmp_file.seek(0, os.SEEK_END) @@ -308,11 +312,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal tmp_file.write(data) tmp_file.flush() - scp_cmd = ['scp'] + self.ssh_cmd + [tmp_file.name, f"{self.username}@{self.host}:{filename}"] + scp_cmd = ['scp'] + scp_args + [tmp_file.name, f"{self.username}@{self.host}:{filename}"] subprocess.run(scp_cmd, check=True) remote_directory = os.path.dirname(filename) - mkdir_cmd = ['ssh'] + self.ssh_cmd + [f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"] + mkdir_cmd = ['ssh'] + self.ssh_args + [f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) @@ -377,7 +381,7 @@ def get_pid(self): return int(self.exec_command("echo $$", encoding=get_default_encoding())) def get_process_children(self, pid): - command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", f"pgrep -P {pid}"] + command = ["ssh"] + self.ssh_args + [f"{self.username}@{self.host}", f"pgrep -P {pid}"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier! Saves Data!