Skip to content

Commit

Permalink
Update establish_ssh_tunnel
Browse files Browse the repository at this point in the history
  • Loading branch information
vshepard committed Jun 6, 2024
1 parent fffb23c commit 1eb9a92
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 5 deletions.
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
readme = f.read()

setup(
version='1.10.1',
version='1.10.2',
name='testgres',
packages=['testgres', 'testgres.operations', 'testgres.helpers'],
description='Testing utility for PostgreSQL and its extensions',
Expand Down
11 changes: 7 additions & 4 deletions testgres/operations/remote_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,11 +76,14 @@ def is_port_open(host, port):
except socket.error:
return False

def establish_ssh_tunnel(self, local_port, remote_port):
def establish_ssh_tunnel(self, local_port, remote_port, host):
"""
Establish an SSH tunnel from a local port to a remote PostgreSQL port.
"""
ssh_cmd = ['-N', '-L', f"{local_port}:localhost:{remote_port}"]
if host != 'localhost':
ssh_cmd = ['-N', '-L', f"localhost:{local_port}:{host}:{remote_port}"]
else:
ssh_cmd = ['-N', '-L', f"{local_port}:{host}:{remote_port}"]
self.tunnel_process = self.exec_command(ssh_cmd, get_process=True, timeout=300)
timeout = 10
start_time = time.time()
Expand Down Expand Up @@ -412,10 +415,10 @@ def db_connect(self, dbname, user, password=None, host="localhost", port=5432):
"""
local_port = reserve_port()
self.tunnel_port = local_port
self.establish_ssh_tunnel(local_port=local_port, remote_port=port)
self.establish_ssh_tunnel(local_port=local_port, remote_port=port, host=host)
try:
conn = pglib.connect(
host=host,
host='localhost',
port=local_port,
database=dbname,
user=user,
Expand Down

0 comments on commit 1eb9a92

Please sign in to comment.