@@ -74,11 +74,11 @@ def is_port_open(host, port):
74
74
except socket .error :
75
75
return False
76
76
77
- def establish_ssh_tunnel (self , local_port , remote_port ):
77
+ def establish_ssh_tunnel (self , local_port , remote_port , host = 'localhost' ):
78
78
"""
79
79
Establish an SSH tunnel from a local port to a remote PostgreSQL port.
80
80
"""
81
- ssh_cmd = ['-N' , '-L' , f"{ local_port } :localhost :{ remote_port } " ]
81
+ ssh_cmd = ['-N' , '-L' , f"{ local_port } :{ host } :{ remote_port } " ]
82
82
self .tunnel_process = self .exec_command (ssh_cmd , get_process = True , timeout = 300 )
83
83
timeout = 10
84
84
start_time = time .time ()
@@ -410,7 +410,7 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
410
410
"""
411
411
local_port = reserve_port ()
412
412
self .tunnel_port = local_port
413
- self .establish_ssh_tunnel (local_port = local_port , remote_port = port )
413
+ self .establish_ssh_tunnel (local_port = local_port , remote_port = port , host = host )
414
414
try :
415
415
conn = pglib .connect (
416
416
host = host ,
@@ -423,8 +423,8 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
423
423
print ("Database connection established successfully." )
424
424
return conn
425
425
except Exception as e :
426
- print (f"Error connecting to the database: { str (e )} " )
426
+ print (f"!!! Error connecting to the database: { str (e )} " )
427
427
if self .tunnel_process :
428
428
self .tunnel_process .terminate ()
429
- print ("SSH tunnel closed due to connection failure." )
429
+ print ("!!! SSH tunnel closed due to connection failure." )
430
430
raise
0 commit comments