Fix ssh command in remote_ops.py

vshepard · vshepard · commit 22d3251b8586 · 2024-05-26T00:41:09.000+02:00
diff --git 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
@@ -1,8 +1,9 @@
-import logging
 import os
+import socket
 import subprocess
 import tempfile
 import platform
+import time
 
 # we support both pg8000 and psycopg2
 try:
@@ -48,10 +49,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
@@ -62,17 +63,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.")
@@ -238,9 +258,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)
 
@@ -283,7 +303,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:
@@ -305,9 +325,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)
@@ -372,7 +392,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)
 
@@ -387,15 +407,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
