Location via proxy:   [ UP ]  
[Report a bug]   [Manage cookies]                
Skip to content

Remove SSH tunnel #136

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions testgres/node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down
11 changes: 10 additions & 1 deletion testgres/operations/os_ops.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import getpass
import locale
import sys

try:
import psycopg2 as pglib # noqa: F401
Expand All @@ -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):
Expand All @@ -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()

Expand Down
55 changes: 9 additions & 46 deletions testgres/operations/remote_ops.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import getpass
import os
import logging
import platform
import subprocess
import tempfile
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
9 changes: 6 additions & 3 deletions testgres/plugins/pg_probackup2/pg_probackup2/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
):
Expand All @@ -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',
Expand Down