@@ -48,10 +48,10 @@ def __init__(self, conn_params: ConnectionParams):
4848self .ssh_key = conn_params .ssh_key
4949self .port = conn_params .port
5050self .ssh_cmd = ["-o StrictHostKeyChecking=no" ]
51- if self .ssh_key :
52- self .ssh_cmd += ["-i" ,self .ssh_key ]
5351if self .port :
5452self .ssh_cmd += ["-p" ,self .port ]
53+ if self .ssh_key :
54+ self .ssh_cmd += ["-i" ,self .ssh_key ]
5555self .remote = True
5656self .username = conn_params .username or self .get_user ()
5757self .tunnel_process = None
@@ -283,6 +283,7 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
283283mode = "r+b" if binary else "r+"
284284
285285with tempfile .NamedTemporaryFile (mode = mode ,delete = False )as tmp_file :
286+ # Because in scp we set up port using -P option instead -p
286287scp_ssh_cmd = ['-P' if x == '-p' else x for x in self .ssh_cmd ]
287288
288289if not truncate :
@@ -302,12 +303,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
302303tmp_file .write (data )
303304
304305tmp_file .flush ()
305- # Because in scp we set up port using -P option
306306scp_cmd = ['scp' ]+ scp_ssh_cmd + [tmp_file .name ,f"{ self .username } { self .host } { filename } ]
307307subprocess .run (scp_cmd ,check = True )
308-
309308remote_directory = os .path .dirname (filename )
310- mkdir_cmd = ['ssh' ]+ scp_ssh_cmd + [f"{ self .username } { self .host } ,f"mkdir -p{ remote_directory } ]
309+
310+ mkdir_cmd = ['ssh' ]+ self .ssh_cmd + [f"{ self .username } { self .host } ,f'mkdir -p{ remote_directory } ]
311311subprocess .run (mkdir_cmd ,check = True )
312312
313313os .remove (tmp_file .name )
@@ -385,9 +385,10 @@ def get_process_children(self, pid):
385385# Database control
386386def db_connect (self ,dbname ,user ,password = None ,host = "localhost" ,port = 5432 ):
387387"""
388- Established SSH tunnel andConnects to a PostgreSQL
388+ Establish SSH tunnel andconnect to a PostgreSQL database.
389389 """
390- self .establish_ssh_tunnel (local_port = reserve_port (),remote_port = 5432 )
390+ self .establish_ssh_tunnel (local_port = port ,remote_port = self .conn_params .port )
391+
391392try :
392393conn = pglib .connect (
393394host = host ,
@@ -396,6 +397,11 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
396397user = user ,
397398password = password ,
398399 )
400+ print ("Database connection established successfully." )
399401return conn
400402except Exception as e :
401- raise Exception (f"Could not connect to the database. Error:{ e } )
403+ print (f"Error connecting to the database:{ str (e )} )
404+ if self .tunnel_process :
405+ self .tunnel_process .terminate ()
406+ print ("SSH tunnel closed due to connection failure." )
407+ raise
