diff --git a/testgres/node.py b/testgres/node.py index 13d13294..479ea4ec 100644 --- a/testgres/node.py +++ b/testgres/node.py @@ -528,7 +528,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 @@ -1671,9 +1673,15 @@ 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()): + 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() + 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 0b5efff9..76284049 100644 --- a/testgres/operations/os_ops.py +++ b/testgres/operations/os_ops.py @@ -1,4 +1,6 @@ +import getpass import locale +import sys try: import psycopg2 as pglib # noqa: F401 @@ -24,7 +26,7 @@ def get_default_encoding(): class OsOperations: def __init__(self, username=None): self.ssh_key = None - self.username = username + self.username = username or getpass.getuser() # Command execution def exec_command(self, cmd, **kwargs): @@ -34,6 +36,13 @@ def exec_command(self, cmd, **kwargs): def environ(self, var_name): raise NotImplementedError() + def cwd(self): + 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 83965336..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 @@ -55,40 +54,10 @@ 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.tunnel_process = None def __enter__(self): return 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) - - def close_ssh_tunnel(self): - if hasattr(self, 'tunnel_process'): - self.tunnel_process.terminate() - self.tunnel_process.wait() - 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): @@ -293,6 +262,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal 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'] + scp_args + [f"{self.ssh_dest}:{filename}", tmp_file.name] subprocess.run(scp_cmd, check=False) # The file might not exist yet @@ -391,18 +361,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 diff --git a/testgres/plugins/pg_probackup2/pg_probackup2/app.py b/testgres/plugins/pg_probackup2/pg_probackup2/app.py index 1a4ca9e7..ffad24d3 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 @@ -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 += ['--remote-proto=ssh', f'--remote-host={remote_host}'] + if remote_port: + cmd_list.append(f'--remote-port={remote_port}') if self.verbose: cmd_list += [ '--log-level-file=VERBOSE',