From e8f4da9f63ac2681d7a10996a3eea625378aea6c Mon Sep 17 00:00:00 2001 From: vshepard Date: Wed, 10 Apr 2024 15:36:06 +0200 Subject: [PATCH 01/27] Add port to connection parameters --- testgres/operations/os_ops.py | 3 ++- testgres/operations/remote_ops.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) 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 f182768b..34b76500 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -47,10 +47,13 @@ def __init__(self, conn_params: ConnectionParams): self.conn_params = conn_params self.host = conn_params.host self.ssh_key = conn_params.ssh_key + self.port = conn_params.port if self.ssh_key: self.ssh_cmd = ["-i", self.ssh_key] else: self.ssh_cmd = [] + if self.port: + self.ssh_cmd = ["-p", self.port] self.remote = True self.username = conn_params.username or self.get_user() self.add_known_host(self.host) From 80ba614ffc50b42949ed28b097b6b79a3d840f31 Mon Sep 17 00:00:00 2001 From: vshepard Date: Thu, 11 Apr 2024 09:13:23 +0200 Subject: [PATCH 02/27] Add probackup_path to init ProbackupApp --- testgres/plugins/pg_probackup2/pg_probackup2/app.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/testgres/plugins/pg_probackup2/pg_probackup2/app.py b/testgres/plugins/pg_probackup2/pg_probackup2/app.py index 1ab71109..b0e6b960 100644 --- a/testgres/plugins/pg_probackup2/pg_probackup2/app.py +++ b/testgres/plugins/pg_probackup2/pg_probackup2/app.py @@ -43,14 +43,14 @@ def __str__(self): class ProbackupApp: def __init__(self, test_class: unittest.TestCase, - pg_node, pb_log_path, test_env, auto_compress_alg, backup_dir): + pg_node, pb_log_path, test_env, auto_compress_alg, backup_dir, probackup_path=None): self.test_class = test_class self.pg_node = pg_node self.pb_log_path = pb_log_path self.test_env = test_env self.auto_compress_alg = auto_compress_alg self.backup_dir = backup_dir - self.probackup_path = init_params.probackup_path + self.probackup_path = probackup_path or init_params.probackup_path self.probackup_old_path = init_params.probackup_old_path self.remote = init_params.remote self.verbose = init_params.verbose From dc775c20df12e5ad4383bbe551cf2964d72c2bd6 Mon Sep 17 00:00:00 2001 From: vshepard Date: Fri, 12 Apr 2024 11:50:44 +0200 Subject: [PATCH 03/27] Add StrictHostKeyChecking=no for ssh connect --- testgres/operations/remote_ops.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 34b76500..eafdaf2a 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -48,15 +48,13 @@ def __init__(self, conn_params: ConnectionParams): self.host = conn_params.host self.ssh_key = conn_params.ssh_key self.port = conn_params.port + self.ssh_cmd = ["-o StrictHostKeyChecking=no"] if self.ssh_key: - self.ssh_cmd = ["-i", self.ssh_key] - else: - self.ssh_cmd = [] + self.ssh_cmd += ["-i", self.ssh_key] if self.port: - self.ssh_cmd = ["-p", self.port] + self.ssh_cmd += ["-p", self.port] self.remote = True self.username = conn_params.username or self.get_user() - self.add_known_host(self.host) self.tunnel_process = None def __enter__(self): @@ -80,16 +78,6 @@ def close_ssh_tunnel(self): else: print("No active tunnel to close.") - def add_known_host(self, host): - known_hosts_path = os.path.expanduser("~/.ssh/known_hosts") - cmd = 'ssh-keyscan -H %s >> %s' % (host, known_hosts_path) - - try: - subprocess.check_call(cmd, shell=True) - logging.info("Successfully added %s to known_hosts." % host) - except subprocess.CalledProcessError as e: - raise Exception("Failed to add %s to known_hosts. Error: %s" % (host, str(e))) - def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=None, timeout=None): From da223402fa95402ce89fcf9bce8ac5df899c7951 Mon Sep 17 00:00:00 2001 From: vshepard Date: Sun, 14 Apr 2024 00:16:49 +0300 Subject: [PATCH 04/27] Add StrictHostKeyChecking=no for ssh connect --- testgres/operations/remote_ops.py | 1 + testgres/utils.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index eafdaf2a..c0b7011e 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -16,6 +16,7 @@ raise ImportError("You must have psycopg2 or pg8000 modules installed") from ..exceptions import ExecUtilException +from ..utils import reserve_port from .os_ops import OsOperations, ConnectionParams, get_default_encoding error_markers = [b'error', b'Permission denied', b'fatal', b'No such file or directory'] diff --git a/testgres/utils.py b/testgres/utils.py index 745a2555..633a9b2f 100644 --- a/testgres/utils.py +++ b/testgres/utils.py @@ -74,6 +74,8 @@ def execute_utility(args, logfile=None, verbose=False): # write new log entry if possible if logfile: + if not tconf.os_ops.path_exists(logfile): + tconf.os_ops.touch(logfile) try: tconf.os_ops.write(filename=logfile, data=args, truncate=True) if out: From d08325bec626616d9cdab9ccf92f81f6cd82eb0a Mon Sep 17 00:00:00 2001 From: vshepard Date: Sun, 14 Apr 2024 01:06:05 +0300 Subject: [PATCH 05/27] Add port to scp command --- testgres/operations/remote_ops.py | 9 ++++++--- testgres/utils.py | 2 -- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index c0b7011e..f5f0a1a5 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -285,8 +285,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: + scp_ssh_cmd = ['-P' if x == '-p' else x for x in self.ssh_cmd] + if not truncate: - scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.username}@{self.host}:{filename}", tmp_file.name] + scp_cmd = ['scp'] + scp_ssh_cmd + [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) @@ -302,11 +304,12 @@ 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}"] + # Because in scp we set up port using -P option + scp_cmd = ['scp'] + scp_ssh_cmd + [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'] + scp_ssh_cmd + [f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) diff --git a/testgres/utils.py b/testgres/utils.py index 633a9b2f..745a2555 100644 --- a/testgres/utils.py +++ b/testgres/utils.py @@ -74,8 +74,6 @@ def execute_utility(args, logfile=None, verbose=False): # write new log entry if possible if logfile: - if not tconf.os_ops.path_exists(logfile): - tconf.os_ops.touch(logfile) try: tconf.os_ops.write(filename=logfile, data=args, truncate=True) if out: From f786b8739d09aebccb2270671841461711d46b40 Mon Sep 17 00:00:00 2001 From: vshepard Date: Thu, 25 Apr 2024 09:12:46 +0200 Subject: [PATCH 06/27] PBCKP-781 change defaukt vakue PG_PROBACKUP_S3_HTTPS on ON --- testgres/operations/remote_ops.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index f5f0a1a5..005d6ac4 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -50,10 +50,10 @@ def __init__(self, conn_params: ConnectionParams): self.ssh_key = conn_params.ssh_key self.port = conn_params.port self.ssh_cmd = ["-o StrictHostKeyChecking=no"] - if self.ssh_key: - self.ssh_cmd += ["-i", self.ssh_key] if self.port: self.ssh_cmd += ["-p", self.port] + if self.ssh_key: + self.ssh_cmd += ["-i", self.ssh_key] self.remote = True self.username = conn_params.username or self.get_user() self.tunnel_process = None @@ -285,6 +285,7 @@ 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: + # Because in scp we set up port using -P option instead -p scp_ssh_cmd = ['-P' if x == '-p' else x for x in self.ssh_cmd] if not truncate: @@ -304,12 +305,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal tmp_file.write(data) tmp_file.flush() - # Because in scp we set up port using -P option scp_cmd = ['scp'] + scp_ssh_cmd + [tmp_file.name, f"{self.username}@{self.host}:{filename}"] subprocess.run(scp_cmd, check=True) - remote_directory = os.path.dirname(filename) - mkdir_cmd = ['ssh'] + scp_ssh_cmd + [f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"] + + mkdir_cmd = ['ssh'] + self.ssh_cmd + [f"{self.username}@{self.host}", f'mkdir -p {remote_directory}'] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) @@ -387,9 +387,10 @@ def get_process_children(self, pid): # Database control def db_connect(self, dbname, user, password=None, host="localhost", port=5432): """ - Established SSH tunnel and Connects to a PostgreSQL + Establish SSH tunnel and connect to a PostgreSQL database. """ - self.establish_ssh_tunnel(local_port=reserve_port(), remote_port=5432) + self.establish_ssh_tunnel(local_port=port, remote_port=self.conn_params.port) + try: conn = pglib.connect( host=host, @@ -398,6 +399,11 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432): user=user, password=password, ) + print("Database connection established successfully.") return conn except Exception as e: - raise Exception(f"Could not connect to the database. Error: {e}") + print(f"Error connecting to the database: {str(e)}") + if self.tunnel_process: + self.tunnel_process.terminate() + print("SSH tunnel closed due to connection failure.") + raise From 48594f67a4fecb6dc2db4620ef143ad3d2348271 Mon Sep 17 00:00:00 2001 From: vshepard Date: Fri, 24 May 2024 22:42:47 +0200 Subject: [PATCH 07/27] Fix ssh command in remote_ops.py --- testgres/node.py | 3 +- testgres/operations/remote_ops.py | 47 ++++++++++++++++++++++--------- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index d1784cb9..66d211dc 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -529,7 +529,8 @@ def get_auth_method(t): u"host\treplication\tall\t127.0.0.1/32\t{}\n".format(auth_host), u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host), u"host\treplication\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host), - u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host) + u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host), + u"host\tall\tall\tall\t{}\n".format(auth_host) ] # yapf: disable # write missing lines diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 005d6ac4..cd9ce201 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -1,8 +1,9 @@ -import logging import os +import socket import subprocess import tempfile import platform +import time from ..utils import reserve_port @@ -50,10 +51,10 @@ def __init__(self, conn_params: ConnectionParams): self.ssh_key = conn_params.ssh_key self.port = conn_params.port self.ssh_cmd = ["-o StrictHostKeyChecking=no"] - if self.port: - self.ssh_cmd += ["-p", self.port] if self.ssh_key: self.ssh_cmd += ["-i", self.ssh_key] + if self.port: + self.ssh_cmd += ["-p", self.port] self.remote = True self.username = conn_params.username or self.get_user() self.tunnel_process = None @@ -64,17 +65,36 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): self.close_ssh_tunnel() + @staticmethod + def is_port_open(host, port): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) # Таймаут для попытки соединения + try: + sock.connect((host, port)) + return True + except socket.error: + return False + def establish_ssh_tunnel(self, local_port, remote_port): """ Establish an SSH tunnel from a local port to a remote PostgreSQL port. """ ssh_cmd = ['-N', '-L', f"{local_port}:localhost:{remote_port}"] self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300) + timeout = 10 + start_time = time.time() + while time.time() - start_time < timeout: + if self.is_port_open('localhost', local_port): + print("SSH tunnel established.") + return + time.sleep(0.5) + raise Exception("Failed to establish SSH tunnel within the timeout period.") def close_ssh_tunnel(self): - if hasattr(self, 'tunnel_process'): + if self.tunnel_process: self.tunnel_process.terminate() self.tunnel_process.wait() + print("SSH tunnel closed.") del self.tunnel_process else: print("No active tunnel to close.") @@ -240,9 +260,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" + f"{self.username}@{self.host}"] + self.ssh_cmd + [f"mktemp -d {prefix}XXXXX"] else: - command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", "mktemp -d"] + command = ["ssh", f"{self.username}@{self.host}"] + self.ssh_cmd + ["mktemp -d"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -285,7 +305,7 @@ 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: - # Because in scp we set up port using -P option instead -p + # Because in scp we set up port using -P option scp_ssh_cmd = ['-P' if x == '-p' else x for x in self.ssh_cmd] if not truncate: @@ -307,9 +327,9 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal tmp_file.flush() scp_cmd = ['scp'] + scp_ssh_cmd + [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}'] + remote_directory = os.path.dirname(filename) + mkdir_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + [f"mkdir -p {remote_directory}"] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) @@ -374,7 +394,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", f"{self.username}@{self.host}"] + self.ssh_cmd + [f"pgrep -P {pid}"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -389,15 +409,16 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432): """ Establish SSH tunnel and connect to a PostgreSQL database. """ - self.establish_ssh_tunnel(local_port=port, remote_port=self.conn_params.port) - + local_port = reserve_port() + self.establish_ssh_tunnel(local_port=local_port, remote_port=port) try: conn = pglib.connect( host=host, - port=port, + port=local_port, database=dbname, user=user, password=password, + timeout=10 ) print("Database connection established successfully.") return conn From fffb23c9b00eac3519230cb8cb9310e130fd7e35 Mon Sep 17 00:00:00 2001 From: vshepard Date: Sun, 2 Jun 2024 23:30:49 +0200 Subject: [PATCH 08/27] Add tunnel_port to remote_ops.py --- testgres/operations/remote_ops.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index cd9ce201..958dc6c2 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -58,6 +58,7 @@ def __init__(self, conn_params: ConnectionParams): self.remote = True self.username = conn_params.username or self.get_user() self.tunnel_process = None + self.tunnel_port = None def __enter__(self): return self @@ -410,6 +411,7 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432): Establish SSH tunnel and connect to a PostgreSQL database. """ local_port = reserve_port() + self.tunnel_port = local_port self.establish_ssh_tunnel(local_port=local_port, remote_port=port) try: conn = pglib.connect( From 1eb9a92fd7fee551fc49939c2048b4a1de571b9f Mon Sep 17 00:00:00 2001 From: vshepard Date: Tue, 4 Jun 2024 20:42:40 +0200 Subject: [PATCH 09/27] Update establish_ssh_tunnel --- setup.py | 2 +- testgres/operations/remote_ops.py | 11 +++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index 412e8823..4f8837c0 100755 --- a/setup.py +++ b/setup.py @@ -27,7 +27,7 @@ readme = f.read() setup( - version='1.10.1', + version='1.10.2', name='testgres', packages=['testgres', 'testgres.operations', 'testgres.helpers'], description='Testing utility for PostgreSQL and its extensions', diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 958dc6c2..12f85403 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -76,11 +76,14 @@ def is_port_open(host, port): except socket.error: return False - def establish_ssh_tunnel(self, local_port, remote_port): + def establish_ssh_tunnel(self, local_port, remote_port, host): """ Establish an SSH tunnel from a local port to a remote PostgreSQL port. """ - ssh_cmd = ['-N', '-L', f"{local_port}:localhost:{remote_port}"] + if host != 'localhost': + ssh_cmd = ['-N', '-L', f"localhost:{local_port}:{host}:{remote_port}"] + else: + ssh_cmd = ['-N', '-L', f"{local_port}:{host}:{remote_port}"] self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300) timeout = 10 start_time = time.time() @@ -412,10 +415,10 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432): """ local_port = reserve_port() self.tunnel_port = local_port - self.establish_ssh_tunnel(local_port=local_port, remote_port=port) + self.establish_ssh_tunnel(local_port=local_port, remote_port=port, host=host) try: conn = pglib.connect( - host=host, + host='localhost', port=local_port, database=dbname, user=user, From 0611d1098fcc97ae1592df9610811a0952f33a61 Mon Sep 17 00:00:00 2001 From: vshepard Date: Fri, 7 Jun 2024 01:01:29 +0200 Subject: [PATCH 10/27] Change pg_hba.conf default params --- testgres/node.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/testgres/node.py b/testgres/node.py index 66d211dc..e94a8e47 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -530,7 +530,8 @@ def get_auth_method(t): u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host), u"host\treplication\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host), u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host), - u"host\tall\tall\tall\t{}\n".format(auth_host) + u"host\tall\tall\tall\t{}\n".format(auth_host), + u"host\treplication\tall\tall\t{}\n".format(auth_host) ] # yapf: disable # write missing lines From 26217241ea7730c41d62bbebd292ef56bff50bf5 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Thu, 13 Jun 2024 09:39:32 +0700 Subject: [PATCH 11/27] Do not use a tunnel --- testgres/operations/remote_ops.py | 52 +++++-------------------------- 1 file changed, 8 insertions(+), 44 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 12f85403..71c4b4b6 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -5,8 +5,6 @@ import platform import time -from ..utils import reserve_port - # we support both pg8000 and psycopg2 try: import psycopg2 as pglib @@ -17,7 +15,6 @@ raise ImportError("You must have psycopg2 or pg8000 modules installed") from ..exceptions import ExecUtilException -from ..utils import reserve_port from .os_ops import OsOperations, ConnectionParams, get_default_encoding error_markers = [b'error', b'Permission denied', b'fatal', b'No such file or directory'] @@ -76,24 +73,6 @@ def is_port_open(host, port): except socket.error: return False - def establish_ssh_tunnel(self, local_port, remote_port, host): - """ - Establish an SSH tunnel from a local port to a remote PostgreSQL port. - """ - if host != 'localhost': - ssh_cmd = ['-N', '-L', f"localhost:{local_port}:{host}:{remote_port}"] - else: - ssh_cmd = ['-N', '-L', f"{local_port}:{host}:{remote_port}"] - self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300) - timeout = 10 - start_time = time.time() - while time.time() - start_time < timeout: - if self.is_port_open('localhost', local_port): - print("SSH tunnel established.") - return - time.sleep(0.5) - raise Exception("Failed to establish SSH tunnel within the timeout period.") - def close_ssh_tunnel(self): if self.tunnel_process: self.tunnel_process.terminate() @@ -410,26 +389,11 @@ def get_process_children(self, pid): # Database control def db_connect(self, dbname, user, password=None, host="localhost", port=5432): - """ - Establish SSH tunnel and connect to a PostgreSQL database. - """ - local_port = reserve_port() - self.tunnel_port = local_port - self.establish_ssh_tunnel(local_port=local_port, remote_port=port, host=host) - try: - conn = pglib.connect( - host='localhost', - port=local_port, - database=dbname, - user=user, - password=password, - timeout=10 - ) - print("Database connection established successfully.") - return conn - except Exception as e: - print(f"Error connecting to the database: {str(e)}") - if self.tunnel_process: - self.tunnel_process.terminate() - print("SSH tunnel closed due to connection failure.") - raise + conn = pglib.connect( + host=host, + port=port, + database=dbname, + user=user, + password=password, + ) + return conn From 08d6b0464e8231a7e5e881d9fb3cad70221f7703 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Thu, 13 Jun 2024 13:56:15 +0700 Subject: [PATCH 12/27] Make use of the default SSH user --- testgres/operations/remote_ops.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 01251e1c..13188053 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -51,6 +51,7 @@ def __init__(self, conn_params: ConnectionParams): self.ssh_cmd = [] self.remote = True self.username = conn_params.username or self.get_user() + self.ssh_dest = f"{self.username}@{self.host}" if self.username else "{self.host}" self.add_known_host(self.host) self.tunnel_process = None @@ -95,9 +96,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', self.ssh_dest] + self.ssh_cmd + [cmd] elif isinstance(cmd, list): - ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + cmd + ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_cmd + cmd process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if get_process: return process @@ -246,9 +247,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_cmd + [self.ssh_dest, f"mktemp -d {prefix}XXXXX"] else: - command = ["ssh"] + self.ssh_cmd + [f"{self.username}@{self.host}", "mktemp -d"] + command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, "mktemp -d"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -292,7 +293,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal with tempfile.NamedTemporaryFile(mode=mode, delete=False) as tmp_file: if not truncate: - scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.username}@{self.host}:{filename}", tmp_file.name] + scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.ssh_dest}:{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 +309,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'] + self.ssh_cmd + [tmp_file.name, f"{self.ssh_dest}:{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_cmd + [self.ssh_dest, f"mkdir -p {remote_directory}"] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) @@ -377,7 +378,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_cmd + [self.ssh_dest, f"pgrep -P {pid}"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) From 8052f39bdde2cd57589a4ad716761d0a0903483b Mon Sep 17 00:00:00 2001 From: asavchkov Date: Thu, 13 Jun 2024 17:20:53 +0700 Subject: [PATCH 13/27] More merge fixes --- testgres/operations/remote_ops.py | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 0bec47f9..1ce01146 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -3,10 +3,6 @@ import subprocess import tempfile import platform -<<<<<<< HEAD -import time -======= ->>>>>>> master # we support both pg8000 and psycopg2 try: @@ -57,11 +53,8 @@ def __init__(self, conn_params: ConnectionParams): self.ssh_cmd += ["-p", self.port] self.remote = True self.username = conn_params.username or self.get_user() -<<<<<<< HEAD -======= self.ssh_dest = f"{self.username}@{self.host}" if self.username else "{self.host}" self.add_known_host(self.host) ->>>>>>> default-ssh-user self.tunnel_process = None self.tunnel_port = None @@ -251,15 +244,9 @@ def mkdtemp(self, prefix=None): - prefix (str): The prefix of the temporary directory name. """ if prefix: -<<<<<<< HEAD - command = ["ssh" + f"{self.username}@{self.host}"] + self.ssh_cmd + [f"mktemp -d {prefix}XXXXX"] - else: - command = ["ssh", f"{self.username}@{self.host}"] + self.ssh_cmd + ["mktemp -d"] -======= command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, f"mktemp -d {prefix}XXXXX"] else: command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, "mktemp -d"] ->>>>>>> default-ssh-user result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -306,11 +293,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal scp_ssh_cmd = ['-P' if x == '-p' else x for x in self.ssh_cmd] if not truncate: -<<<<<<< HEAD - scp_cmd = ['scp'] + scp_ssh_cmd + [f"{self.username}@{self.host}:{filename}", tmp_file.name] -======= scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.ssh_dest}:{filename}", tmp_file.name] ->>>>>>> default-ssh-user subprocess.run(scp_cmd, check=False) # The file might not exist yet tmp_file.seek(0, os.SEEK_END) @@ -326,19 +309,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal tmp_file.write(data) tmp_file.flush() -<<<<<<< HEAD - scp_cmd = ['scp'] + scp_ssh_cmd + [tmp_file.name, f"{self.username}@{self.host}:{filename}"] - subprocess.run(scp_cmd, check=True) - - remote_directory = os.path.dirname(filename) - mkdir_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + [f"mkdir -p {remote_directory}"] -======= scp_cmd = ['scp'] + self.ssh_cmd + [tmp_file.name, f"{self.ssh_dest}:{filename}"] subprocess.run(scp_cmd, check=True) remote_directory = os.path.dirname(filename) mkdir_cmd = ['ssh'] + self.ssh_cmd + [self.ssh_dest, f"mkdir -p {remote_directory}"] ->>>>>>> default-ssh-user subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) From c23695954d8ddae62cd9b8f6eee77f1504f96e6b Mon Sep 17 00:00:00 2001 From: asavchkov Date: Thu, 13 Jun 2024 21:17:00 +0700 Subject: [PATCH 14/27] Do not get_user --- testgres/operations/remote_ops.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 1ce01146..3aef00a5 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -52,8 +52,8 @@ def __init__(self, conn_params: ConnectionParams): if self.port: self.ssh_cmd += ["-p", self.port] self.remote = True - self.username = conn_params.username or self.get_user() - self.ssh_dest = f"{self.username}@{self.host}" if self.username else "{self.host}" + self.username = conn_params.username + self.ssh_dest = f"{self.username}@{self.host}" if self.username else self.host self.add_known_host(self.host) self.tunnel_process = None self.tunnel_port = None @@ -170,10 +170,6 @@ def set_env(self, var_name: str, var_val: str): """ return self.exec_command("export {}={}".format(var_name, var_val)) - # Get environment variables - def get_user(self): - return self.exec_command("echo $USER", encoding=get_default_encoding()).strip() - def get_name(self): cmd = 'python3 -c "import os; print(os.name)"' return self.exec_command(cmd, encoding=get_default_encoding()).strip() From 6b4619e7c71250e6a403ab0e0be5e73345555255 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Thu, 13 Jun 2024 21:27:00 +0700 Subject: [PATCH 15/27] Remove get_user --- testgres/operations/remote_ops.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 13188053..b57627cc 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -50,8 +50,8 @@ def __init__(self, conn_params: ConnectionParams): else: self.ssh_cmd = [] self.remote = True - self.username = conn_params.username or self.get_user() - self.ssh_dest = f"{self.username}@{self.host}" if self.username else "{self.host}" + self.username = conn_params.username + self.ssh_dest = f"{self.username}@{self.host}" if self.username else self.host self.add_known_host(self.host) self.tunnel_process = None @@ -173,10 +173,6 @@ def set_env(self, var_name: str, var_val: str): """ return self.exec_command("export {}={}".format(var_name, var_val)) - # Get environment variables - def get_user(self): - return self.exec_command("echo $USER", encoding=get_default_encoding()).strip() - def get_name(self): cmd = 'python3 -c "import os; print(os.name)"' return self.exec_command(cmd, encoding=get_default_encoding()).strip() From 00f7fd54a2e646b97689efd30de298f3dc211ae0 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Fri, 14 Jun 2024 09:40:19 +0700 Subject: [PATCH 16/27] Fix the arg variable name --- testgres/operations/remote_ops.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index b57627cc..682a4900 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -46,9 +46,9 @@ def __init__(self, conn_params: ConnectionParams): self.host = conn_params.host self.ssh_key = conn_params.ssh_key if self.ssh_key: - self.ssh_cmd = ["-i", self.ssh_key] + self.ssh_args = ["-i", self.ssh_key] else: - self.ssh_cmd = [] + self.ssh_args = [] self.remote = True self.username = conn_params.username self.ssh_dest = f"{self.username}@{self.host}" if self.username else self.host @@ -96,9 +96,9 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, """ ssh_cmd = [] if isinstance(cmd, str): - ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_cmd + [cmd] + ssh_cmd = ['ssh'] + self.ssh_args + [self.ssh_dest + cmd] elif isinstance(cmd, list): - ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_cmd + cmd + ssh_cmd = ['ssh'] + self.ssh_args + [self.ssh_dest] + cmd process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if get_process: return process @@ -243,9 +243,9 @@ def mkdtemp(self, prefix=None): - prefix (str): The prefix of the temporary directory name. """ if prefix: - command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, f"mktemp -d {prefix}XXXXX"] + command = ["ssh"] + self.ssh_args + [self.ssh_dest, f"mktemp -d {prefix}XXXXX"] else: - command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, "mktemp -d"] + command = ["ssh"] + self.ssh_args + [self.ssh_dest, "mktemp -d"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -289,7 +289,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal with tempfile.NamedTemporaryFile(mode=mode, delete=False) as tmp_file: if not truncate: - scp_cmd = ['scp'] + self.ssh_cmd + [f"{self.ssh_dest}:{filename}", tmp_file.name] + scp_cmd = ['scp'] + self.ssh_args + [f"{self.ssh_dest}:{filename}", tmp_file.name] subprocess.run(scp_cmd, check=False) # The file might not exist yet tmp_file.seek(0, os.SEEK_END) @@ -305,11 +305,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.ssh_dest}:{filename}"] + scp_cmd = ['scp'] + self.ssh_args + [tmp_file.name, f"{self.ssh_dest}:{filename}"] subprocess.run(scp_cmd, check=True) remote_directory = os.path.dirname(filename) - mkdir_cmd = ['ssh'] + self.ssh_cmd + [self.ssh_dest, f"mkdir -p {remote_directory}"] + mkdir_cmd = ['ssh'] + self.ssh_args + [self.ssh_dest, f"mkdir -p {remote_directory}"] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) @@ -374,7 +374,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 + [self.ssh_dest, f"pgrep -P {pid}"] + command = ["ssh"] + self.ssh_args + [self.ssh_dest, f"pgrep -P {pid}"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) From 19e9436c2f32e29a93d07e768046a6f000f7b494 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Fri, 14 Jun 2024 09:40:19 +0700 Subject: [PATCH 17/27] Fix the arg variable name --- testgres/operations/remote_ops.py | 81 ++++++++++++++----------------- 1 file changed, 37 insertions(+), 44 deletions(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index b57627cc..7c774a1c 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -1,5 +1,5 @@ -import logging import os +import socket import subprocess import tempfile import platform @@ -45,15 +45,18 @@ def __init__(self, conn_params: ConnectionParams): self.conn_params = conn_params self.host = conn_params.host self.ssh_key = conn_params.ssh_key + self.port = conn_params.port + 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 self.ssh_dest = f"{self.username}@{self.host}" if self.username else self.host self.add_known_host(self.host) self.tunnel_process = None + self.tunnel_port = None def __enter__(self): return self @@ -61,31 +64,25 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): self.close_ssh_tunnel() - def establish_ssh_tunnel(self, local_port, remote_port): - """ - Establish an SSH tunnel from a local port to a remote PostgreSQL port. - """ - ssh_cmd = ['-N', '-L', f"{local_port}:localhost:{remote_port}"] - self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300) + @staticmethod + def is_port_open(host, port): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(1) # Таймаут для попытки соединения + try: + sock.connect((host, port)) + return True + except socket.error: + return False def close_ssh_tunnel(self): - if hasattr(self, 'tunnel_process'): + if self.tunnel_process: self.tunnel_process.terminate() self.tunnel_process.wait() + print("SSH tunnel closed.") del self.tunnel_process else: print("No active tunnel to close.") - def add_known_host(self, host): - known_hosts_path = os.path.expanduser("~/.ssh/known_hosts") - cmd = 'ssh-keyscan -H %s >> %s' % (host, known_hosts_path) - - try: - subprocess.check_call(cmd, shell=True) - logging.info("Successfully added %s to known_hosts." % host) - except subprocess.CalledProcessError as e: - raise Exception("Failed to add %s to known_hosts. Error: %s" % (host, str(e))) - def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, encoding=None, shell=True, text=False, input=None, stdin=None, stdout=None, stderr=None, get_process=None, timeout=None): @@ -96,9 +93,9 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False, """ ssh_cmd = [] if isinstance(cmd, str): - ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_cmd + [cmd] + ssh_cmd = ['ssh'] + self.ssh_args + [self.ssh_dest, cmd] elif isinstance(cmd, list): - ssh_cmd = ['ssh', self.ssh_dest] + self.ssh_cmd + cmd + ssh_cmd = ['ssh'] + self.ssh_args + [self.ssh_dest] + cmd process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) if get_process: return process @@ -243,9 +240,9 @@ def mkdtemp(self, prefix=None): - prefix (str): The prefix of the temporary directory name. """ if prefix: - command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, f"mktemp -d {prefix}XXXXX"] + command = ["ssh"] + self.ssh_args + [self.ssh_dest, f"mktemp -d {prefix}XXXXX"] else: - command = ["ssh"] + self.ssh_cmd + [self.ssh_dest, "mktemp -d"] + command = ["ssh"] + self.ssh_args + [self.ssh_dest, "mktemp -d"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -288,8 +285,11 @@ 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: + # Because in scp we set up port using -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.ssh_dest}:{filename}", tmp_file.name] + scp_cmd = ['scp'] + scp_args + [f"{self.ssh_dest}:{filename}", tmp_file.name] subprocess.run(scp_cmd, check=False) # The file might not exist yet tmp_file.seek(0, os.SEEK_END) @@ -305,11 +305,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.ssh_dest}:{filename}"] + scp_cmd = ['scp'] + scp_args + [tmp_file.name, f"{self.ssh_dest}:{filename}"] subprocess.run(scp_cmd, check=True) remote_directory = os.path.dirname(filename) - mkdir_cmd = ['ssh'] + self.ssh_cmd + [self.ssh_dest, f"mkdir -p {remote_directory}"] + mkdir_cmd = ['ssh'] + self.ssh_args + [self.ssh_dest, f"mkdir -p {remote_directory}"] subprocess.run(mkdir_cmd, check=True) os.remove(tmp_file.name) @@ -374,7 +374,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 + [self.ssh_dest, f"pgrep -P {pid}"] + command = ["ssh"] + self.ssh_args + [self.ssh_dest, f"pgrep -P {pid}"] result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) @@ -386,18 +386,11 @@ def get_process_children(self, pid): # Database control def db_connect(self, dbname, user, password=None, host="localhost", port=5432): - """ - Established SSH tunnel and Connects to a PostgreSQL - """ - self.establish_ssh_tunnel(local_port=port, remote_port=5432) - try: - conn = pglib.connect( - host=host, - port=port, - database=dbname, - user=user, - password=password, - ) - return conn - except Exception as e: - raise Exception(f"Could not connect to the database. Error: {e}") + conn = pglib.connect( + host=host, + port=port, + database=dbname, + user=user, + password=password, + ) + return conn From bc9a62e23258fa09988b66b0f1c12f2ef40c0a8c Mon Sep 17 00:00:00 2001 From: asavchkov Date: Sat, 15 Jun 2024 10:52:09 +0700 Subject: [PATCH 18/27] Implement get_user --- testgres/node.py | 30 ++++++++++-------------------- testgres/operations/local_ops.py | 4 ---- testgres/operations/os_ops.py | 4 ++-- 3 files changed, 12 insertions(+), 26 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index e5e8fd5f..01432446 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -156,9 +156,10 @@ def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionP else: self.os_ops = LocalOperations(conn_params) - self.port = port or reserve_port() - self.host = self.os_ops.host + self.port = port or reserve_port() + # Default node username + self.username = default_username() self.ssh_key = self.os_ops.ssh_key # defaults for __exit__() @@ -683,8 +684,6 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem If False, waits for the instance to be in primary mode. Default is False. max_attempts: """ - if not username: - username = default_username() self.start() if replica: @@ -694,7 +693,7 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem # Call poll_query_until until the expected value is returned self.poll_query_until(query=query, dbname=dbname, - username=username, + username=username or self.username, suppress={InternalError, QueryException, ProgrammingError, @@ -967,15 +966,13 @@ def psql(self, >>> psql(query='select 3', ON_ERROR_STOP=1) """ - # Set default arguments dbname = dbname or default_dbname() - username = username or default_username() psql_params = [ self._get_bin_path("psql"), "-p", str(self.port), "-h", self.host, - "-U", username, + "-U", username or self.username, "-X", # no .psqlrc "-A", # unaligned output "-t", # print rows only @@ -1087,9 +1084,6 @@ def tmpfile(): fname = self.os_ops.mkstemp(prefix=TMP_DUMP) return fname - # Set default arguments - dbname = dbname or default_dbname() - username = username or default_username() filename = filename or tmpfile() _params = [ @@ -1097,8 +1091,8 @@ def tmpfile(): "-p", str(self.port), "-h", self.host, "-f", filename, - "-U", username, - "-d", dbname, + "-U", username or self.username, + "-d", dbname or default_dbname(), "-F", format.value ] # yapf: disable @@ -1118,7 +1112,7 @@ def restore(self, filename, dbname=None, username=None): # Set default arguments dbname = dbname or default_dbname() - username = username or default_username() + username = username or self.username _params = [ self._get_bin_path("pg_restore"), @@ -1388,15 +1382,13 @@ def pgbench(self, if options is None: options = [] - # Set default arguments dbname = dbname or default_dbname() - username = username or default_username() _params = [ self._get_bin_path("pgbench"), "-p", str(self.port), "-h", self.host, - "-U", username, + "-U", username or self.username ] + options # yapf: disable # should be the last one @@ -1463,15 +1455,13 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs): >>> pgbench_run(time=10) """ - # Set default arguments dbname = dbname or default_dbname() - username = username or default_username() _params = [ self._get_bin_path("pgbench"), "-p", str(self.port), "-h", self.host, - "-U", username, + "-U", username or self.username ] + options # yapf: disable for key, value in iteritems(kwargs): diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py index ef360d3b..fc9200a9 100644 --- a/testgres/operations/local_ops.py +++ b/testgres/operations/local_ops.py @@ -130,10 +130,6 @@ def set_env(self, var_name, var_val): # Check if the directory is already in PATH os.environ[var_name] = var_val - # Get environment variables - def get_user(self): - return self.username or getpass.getuser() - def get_name(self): return os.name diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py index dd6613cf..58abdda4 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -1,3 +1,4 @@ +import getpass import locale try: @@ -44,9 +45,8 @@ def set_env(self, var_name, var_val): # Check if the directory is already in PATH raise NotImplementedError() - # Get environment variables def get_user(self): - raise NotImplementedError() + return getpass.getuser() def get_name(self): raise NotImplementedError() From 154d000761e1f20d4f47ed6d82e531b76c345953 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Sat, 15 Jun 2024 22:17:09 +0700 Subject: [PATCH 19/27] Refactor the node username --- testgres/node.py | 3 +-- testgres/operations/local_ops.py | 1 - testgres/operations/os_ops.py | 2 +- testgres/operations/remote_ops.py | 3 +-- 4 files changed, 3 insertions(+), 6 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index 01432446..8736654a 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -158,8 +158,7 @@ def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionP self.host = self.os_ops.host self.port = port or reserve_port() - # Default node username - self.username = default_username() + self.username = self.os_ops.username self.ssh_key = self.os_ops.ssh_key # defaults for __exit__() diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py index fc9200a9..de89e6b1 100644 --- a/testgres/operations/local_ops.py +++ b/testgres/operations/local_ops.py @@ -38,7 +38,6 @@ def __init__(self, conn_params=None): self.host = conn_params.host self.ssh_key = None self.remote = False - self.username = conn_params.username or self.get_user() @staticmethod def _raise_exec_exception(message, command, exit_code, output): diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py index 58abdda4..d2ff255e 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -24,7 +24,7 @@ def get_default_encoding(): class OsOperations: def __init__(self, username=None): self.ssh_key = None - self.username = username + self.username = username or self.get_user() # Command execution def exec_command(self, cmd, **kwargs): diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 7c774a1c..8ecf91d0 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -52,8 +52,7 @@ def __init__(self, conn_params: ConnectionParams): if self.port: self.ssh_args += ["-p", self.port] self.remote = True - self.username = conn_params.username - self.ssh_dest = f"{self.username}@{self.host}" if self.username else self.host + self.ssh_dest = f"{self.username}@{self.host}" if conn_params.username else self.host self.add_known_host(self.host) self.tunnel_process = None self.tunnel_port = None From e6fc5f132762b5730076d657207860c8520266fe Mon Sep 17 00:00:00 2001 From: asavchkov Date: Sat, 15 Jun 2024 23:02:28 +0700 Subject: [PATCH 20/27] Override self.username in both operations --- testgres/node.py | 14 ++++++-------- testgres/operations/local_ops.py | 1 + testgres/operations/os_ops.py | 4 ++-- testgres/operations/remote_ops.py | 2 ++ 4 files changed, 11 insertions(+), 10 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index 8736654a..fe1ff0df 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -63,7 +63,6 @@ from .defaults import \ default_dbname, \ - default_username, \ generate_app_name from .exceptions import \ @@ -158,7 +157,6 @@ def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionP self.host = self.os_ops.host self.port = port or reserve_port() - self.username = self.os_ops.username self.ssh_key = self.os_ops.ssh_key # defaults for __exit__() @@ -692,7 +690,7 @@ def slow_start(self, replica=False, dbname='template1', username=None, max_attem # Call poll_query_until until the expected value is returned self.poll_query_until(query=query, dbname=dbname, - username=username or self.username, + username=username or self.os_ops.username, suppress={InternalError, QueryException, ProgrammingError, @@ -971,7 +969,7 @@ def psql(self, self._get_bin_path("psql"), "-p", str(self.port), "-h", self.host, - "-U", username or self.username, + "-U", username or self.os_ops.username, "-X", # no .psqlrc "-A", # unaligned output "-t", # print rows only @@ -1090,7 +1088,7 @@ def tmpfile(): "-p", str(self.port), "-h", self.host, "-f", filename, - "-U", username or self.username, + "-U", username or self.os_ops.username, "-d", dbname or default_dbname(), "-F", format.value ] # yapf: disable @@ -1111,7 +1109,7 @@ def restore(self, filename, dbname=None, username=None): # Set default arguments dbname = dbname or default_dbname() - username = username or self.username + username = username or self.os_ops.username _params = [ self._get_bin_path("pg_restore"), @@ -1387,7 +1385,7 @@ def pgbench(self, self._get_bin_path("pgbench"), "-p", str(self.port), "-h", self.host, - "-U", username or self.username + "-U", username or self.os_ops.username ] + options # yapf: disable # should be the last one @@ -1460,7 +1458,7 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs): self._get_bin_path("pgbench"), "-p", str(self.port), "-h", self.host, - "-U", username or self.username + "-U", username or self.os_ops.username ] + options # yapf: disable for key, value in iteritems(kwargs): diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py index de89e6b1..313d7060 100644 --- a/testgres/operations/local_ops.py +++ b/testgres/operations/local_ops.py @@ -38,6 +38,7 @@ def __init__(self, conn_params=None): self.host = conn_params.host self.ssh_key = None self.remote = False + self.username = conn_params.username or getpass.getuser() @staticmethod def _raise_exec_exception(message, command, exit_code, output): diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py index d2ff255e..7a4df706 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -24,7 +24,7 @@ def get_default_encoding(): class OsOperations: def __init__(self, username=None): self.ssh_key = None - self.username = username or self.get_user() + self.username = username or getpass.getuser() # Command execution def exec_command(self, cmd, **kwargs): @@ -46,7 +46,7 @@ def set_env(self, var_name, var_val): raise NotImplementedError() def get_user(self): - return getpass.getuser() + return self.username def get_name(self): raise NotImplementedError() diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 8ecf91d0..3954b318 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -1,3 +1,4 @@ +import getpass import os import socket import subprocess @@ -52,6 +53,7 @@ def __init__(self, conn_params: ConnectionParams): if self.port: self.ssh_args += ["-p", self.port] self.remote = True + self.username = conn_params.username or getpass.getuser() self.ssh_dest = f"{self.username}@{self.host}" if conn_params.username else self.host self.add_known_host(self.host) self.tunnel_process = None From 0a0b0cf30ac43df2970d7924f501401a38ac382e Mon Sep 17 00:00:00 2001 From: asavchkov Date: Mon, 17 Jun 2024 18:34:15 +0700 Subject: [PATCH 21/27] Add a remote port parameter to the catchup --- testgres/plugins/pg_probackup2/pg_probackup2/app.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/testgres/plugins/pg_probackup2/pg_probackup2/app.py b/testgres/plugins/pg_probackup2/pg_probackup2/app.py index 1a4ca9e7..d5e7a412 100644 --- a/testgres/plugins/pg_probackup2/pg_probackup2/app.py +++ b/testgres/plugins/pg_probackup2/pg_probackup2/app.py @@ -388,6 +388,7 @@ def catchup_node( backup_mode, source_pgdata, destination_node, options=None, remote_host='localhost', + remote_port=None, expect_error=False, gdb=False ): @@ -401,7 +402,9 @@ def catchup_node( '--destination-pgdata={0}'.format(destination_node.data_dir) ] if self.remote: - cmd_list += ['--remote-proto=ssh', '--remote-host=%s' % remote_host] + cmd_list += [f'--remote-proto=ssh', '--remote-host={remote_host}'] + if remote_port: + cmd_list.append(f'--remote-port={remote_port}') if self.verbose: cmd_list += [ '--log-level-file=VERBOSE', From 1c230162232eefbef008b01009622846142ecad6 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Fri, 21 Jun 2024 11:47:28 +0700 Subject: [PATCH 22/27] Remove add hosts --- testgres/operations/remote_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 3954b318..0576319e 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -55,7 +55,7 @@ def __init__(self, conn_params: ConnectionParams): self.remote = True self.username = conn_params.username or getpass.getuser() self.ssh_dest = f"{self.username}@{self.host}" if conn_params.username else self.host - self.add_known_host(self.host) + #self.add_known_host(self.host) self.tunnel_process = None self.tunnel_port = None From 231c5da22363a49a49de617f0d1b891efeec2d55 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Fri, 21 Jun 2024 12:13:47 +0700 Subject: [PATCH 23/27] Determine test_path if not specified --- testgres/node.py | 14 +++++++++++--- testgres/operations/os_ops.py | 6 ++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index 90821267..d91784da 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -1672,9 +1672,17 @@ def _get_bin_path(self, filename): class NodeApp: - def __init__(self, test_path, nodes_to_cleanup, os_ops=LocalOperations()): - self.test_path = test_path - self.nodes_to_cleanup = nodes_to_cleanup + def __init__(self, test_path=None, nodes_to_cleanup=None, os_ops=LocalOperations()): + print('ALEXEY in nodeapp init', test_path) + if test_path: + if os.path.isabs(test_path): + self.test_path = test_path + else: + self.test_path = os.path.join(os_ops.cwd(), test_path) + else: + self.test_path = os_ops.cwd() + print('ALEXEY in nodeapp resulting test path', self.test_path) + self.nodes_to_cleanup = nodes_to_cleanup if nodes_to_cleanup else [] self.os_ops = os_ops def make_empty( diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py index 95926b36..d172e5d0 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -1,5 +1,6 @@ import getpass import locale +import sys try: import psycopg2 as pglib # noqa: F401 @@ -35,6 +36,11 @@ def exec_command(self, cmd, **kwargs): def environ(self, var_name): raise NotImplementedError() + def cwd(self): + if sys.platform == 'win32': + raise NotImplementedError() + return self.exec_command('pwd').decode().rstrip() + def find_executable(self, executable): raise NotImplementedError() From 14080dda4e8f594f97702b3d22438dac97e51011 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Fri, 21 Jun 2024 23:07:14 +0700 Subject: [PATCH 24/27] Pass DB port to NodeApp --- testgres/node.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/testgres/node.py b/testgres/node.py index fe1ff0df..321ab5a8 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -1677,12 +1677,13 @@ def __init__(self, test_path, nodes_to_cleanup, os_ops=LocalOperations()): def make_empty( self, + port=None, base_dir=None): real_base_dir = os.path.join(self.test_path, base_dir) self.os_ops.rmdirs(real_base_dir, ignore_errors=True) self.os_ops.makedirs(real_base_dir) - node = PostgresNode(base_dir=real_base_dir) + node = PostgresNode(port=port, base_dir=real_base_dir) node.should_rm_dirs = True self.nodes_to_cleanup.append(node) @@ -1690,6 +1691,7 @@ def make_empty( def make_simple( self, + port=None, base_dir=None, set_replication=False, ptrack_enable=False, From f5310888308752aea4d974c844d998c652c6b873 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Fri, 21 Jun 2024 23:42:23 +0700 Subject: [PATCH 25/27] Reorder args --- testgres/node.py | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index 321ab5a8..cb19d4da 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -126,7 +126,7 @@ def __repr__(self): class PostgresNode(object): - def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionParams = ConnectionParams(), bin_dir=None, prefix=None): + def __init__(self, name=None, base_dir=None, port=None, conn_params: ConnectionParams = ConnectionParams(), bin_dir=None, prefix=None): """ PostgresNode constructor. @@ -527,7 +527,9 @@ def get_auth_method(t): u"host\treplication\tall\t127.0.0.1/32\t{}\n".format(auth_host), u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host), u"host\treplication\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host), - u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host) + u"host\tall\tall\t{}/24\t\t{}\n".format(subnet_base, auth_host), + u"host\tall\tall\tall\t{}\n".format(auth_host), + u"host\treplication\tall\tall\t{}\n".format(auth_host) ] # yapf: disable # write missing lines @@ -1670,20 +1672,28 @@ def _get_bin_path(self, filename): class NodeApp: - def __init__(self, test_path, nodes_to_cleanup, os_ops=LocalOperations()): - self.test_path = test_path - self.nodes_to_cleanup = nodes_to_cleanup + def __init__(self, test_path=None, nodes_to_cleanup=None, os_ops=LocalOperations()): + print('ALEXEY in nodeapp init', test_path) + if test_path: + if os.path.isabs(test_path): + self.test_path = test_path + else: + self.test_path = os.path.join(os_ops.cwd(), test_path) + else: + self.test_path = os_ops.cwd() + print('ALEXEY in nodeapp resulting test path', self.test_path) + self.nodes_to_cleanup = nodes_to_cleanup if nodes_to_cleanup else [] self.os_ops = os_ops def make_empty( self, - port=None, - base_dir=None): + base_dir=None, + port=None): real_base_dir = os.path.join(self.test_path, base_dir) self.os_ops.rmdirs(real_base_dir, ignore_errors=True) self.os_ops.makedirs(real_base_dir) - node = PostgresNode(port=port, base_dir=real_base_dir) + node = PostgresNode(base_dir=real_base_dir, port=port) node.should_rm_dirs = True self.nodes_to_cleanup.append(node) @@ -1691,8 +1701,8 @@ def make_empty( def make_simple( self, - port=None, base_dir=None, + port=None, set_replication=False, ptrack_enable=False, initdb_params=[], @@ -1700,7 +1710,7 @@ def make_simple( checksum=True): if checksum and '--data-checksums' not in initdb_params: initdb_params.append('--data-checksums') - node = self.make_empty(base_dir) + node = self.make_empty(base_dir, port) node.init( initdb_params=initdb_params, allow_streaming=set_replication) From 8a38d7cc8691292a869babfb2a3d92b69556e931 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Sat, 22 Jun 2024 08:23:28 +0700 Subject: [PATCH 26/27] Fix an f-string --- testgres/plugins/pg_probackup2/pg_probackup2/app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testgres/plugins/pg_probackup2/pg_probackup2/app.py b/testgres/plugins/pg_probackup2/pg_probackup2/app.py index d5e7a412..daa3ba24 100644 --- a/testgres/plugins/pg_probackup2/pg_probackup2/app.py +++ b/testgres/plugins/pg_probackup2/pg_probackup2/app.py @@ -402,7 +402,7 @@ def catchup_node( '--destination-pgdata={0}'.format(destination_node.data_dir) ] if self.remote: - cmd_list += [f'--remote-proto=ssh', '--remote-host={remote_host}'] + cmd_list += ['--remote-proto=ssh', f'--remote-host={remote_host}'] if remote_port: cmd_list.append(f'--remote-port={remote_port}') if self.verbose: From 14b4402d878db095771db8060ce3b5d0f3335d00 Mon Sep 17 00:00:00 2001 From: asavchkov Date: Sun, 30 Jun 2024 21:10:04 +0700 Subject: [PATCH 27/27] Cwd on Windows --- testgres/node.py | 2 -- testgres/operations/os_ops.py | 8 +++++--- testgres/operations/remote_ops.py | 1 - 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/testgres/node.py b/testgres/node.py index 49489aac..479ea4ec 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -1674,7 +1674,6 @@ def _get_bin_path(self, filename): class NodeApp: def __init__(self, test_path=None, nodes_to_cleanup=None, os_ops=LocalOperations()): - print('ALEXEY in nodeapp init', test_path) if test_path: if os.path.isabs(test_path): self.test_path = test_path @@ -1682,7 +1681,6 @@ def __init__(self, test_path=None, nodes_to_cleanup=None, os_ops=LocalOperations self.test_path = os.path.join(os_ops.cwd(), test_path) else: self.test_path = os_ops.cwd() - print('ALEXEY in nodeapp resulting test path', self.test_path) self.nodes_to_cleanup = nodes_to_cleanup if nodes_to_cleanup else [] self.os_ops = os_ops diff --git a/testgres/operations/os_ops.py b/testgres/operations/os_ops.py index d172e5d0..76284049 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -37,9 +37,11 @@ def environ(self, var_name): raise NotImplementedError() def cwd(self): - if sys.platform == 'win32': - raise NotImplementedError() - return self.exec_command('pwd').decode().rstrip() + if sys.platform == 'linux': + cmd = 'pwd' + elif sys.platform == 'win32': + cmd = 'cd' + return self.exec_command(cmd).decode().rstrip() def find_executable(self, executable): raise NotImplementedError() diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py index 43e506a0..fa031075 100644 --- a/testgres/operations/remote_ops.py +++ b/testgres/operations/remote_ops.py @@ -1,6 +1,5 @@ import getpass import os -import logging import platform import subprocess import tempfile pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy