Remove SSH tunnel (#136)

asavchkov · web-flow · commit a128b12cb18f · 2024-07-01T10:47:40.000+02:00
diff --git 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
@@ -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
@@ -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
@@ -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',
