PBCKP-152 test multihost

Author: v.shepard

testgres/cache.py

@@ -21,17 +21,17 @@
     execute_utility
 
 
-def cached_initdb(data_dir, logfile=None, hostname='localhost', ssh_key=None, params=None):
+def cached_initdb(data_dir, logfile=None, host='localhost', ssh_key=None, user='dev', params=None):
     """
     Perform initdb or use cached node files.
     """
 
     def call_initdb(initdb_dir, log=None):
         try:
-            _params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
+            _params = [get_bin_path("initdb", host, ssh_key), "-D", initdb_dir, "-N"]
 
             # DDD return executions code
-            execute_utility(_params + (params or []), log, hostname=hostname, ssh_key=ssh_key)
+            execute_utility(_params + (params or []), log, host=host, ssh_key=ssh_key, user=user)
         except ExecUtilException as e:
             raise_from(InitNodeException("Failed to run initdb"), e)
 

testgres/node.py

@@ -6,6 +6,8 @@
 import subprocess
 import time
 
+from os_ops import OsOperations
+
 try:
     from collections.abc import Iterable
 except ImportError:
@@ -110,7 +112,7 @@ def __repr__(self):
 
 class PostgresNode(object):
     def __init__(self, name=None, port=None, base_dir=None,
-                 host='127.0.0.1', hostname='locahost', ssh_key=None):
+                 host='locahost', ssh_key=None, username='dev'):
         """
         PostgresNode constructor.
 
@@ -129,10 +131,10 @@ def __init__(self, name=None, port=None, base_dir=None,
 
         # basic
         self.host = host
-        self.hostname = hostname
         self.ssh_key = ssh_key
         self.name = name or generate_app_name()
         self.port = port or reserve_port()
+        self.os_ops = OsOperations(host, ssh_key, username)
 
         # defaults for __exit__()
         self.cleanup_on_good_exit = testgres_config.node_cleanup_on_good_exit
@@ -255,20 +257,14 @@ def base_dir(self):
         if not self._base_dir:
             self._base_dir = mkdtemp(prefix=TMP_NODE)
 
-        # NOTE: it's safe to create a new dir
-        if not os.path.exists(self._base_dir):
-            os.makedirs(self._base_dir)
+        self.os_ops.makedirs(self._base_dir)
 
         return self._base_dir
 
     @property
     def logs_dir(self):
         path = os.path.join(self.base_dir, LOGS_DIR)
-
-        # NOTE: it's safe to create a new dir
-        if not os.path.exists(path):
-            os.makedirs(path)
-
+        self.os_ops.makedirs(path)
         return path
 
     @property
@@ -420,7 +416,7 @@ def _collect_special_files(self):
 
         return result
 
-    def init(self, initdb_params=None, hostname='localhost', ssh_key=None, **kwargs):
+    def init(self, initdb_params=None, host='localhost', ssh_key=None, **kwargs):
         """
         Perform initdb for this node.
 
@@ -429,6 +425,8 @@ def init(self, initdb_params=None, hostname='localhost', ssh_key=None, **kwargs)
             fsync: should this node use fsync to keep data safe?
             unix_sockets: should we enable UNIX sockets?
             allow_streaming: should this node add a hba entry for replication?
+            host: for remote connection
+            ssh_key: for remote connection
 
         Returns:
             This instance of :class:`.PostgresNode`
@@ -437,7 +435,7 @@ def init(self, initdb_params=None, hostname='localhost', ssh_key=None, **kwargs)
         # initialize this PostgreSQL node
         cached_initdb(data_dir=self.data_dir,
                       logfile=self.utils_log_file,
-                      hostname=hostname,
+                      host=host,
                       ssh_key=ssh_key,
                       params=initdb_params)
 
@@ -470,6 +468,7 @@ def default_conf(self,
         hba_conf = os.path.join(self.data_dir, HBA_CONF_FILE)
 
         # filter lines in hba file
+        lines = self.os_ops.readlines(hba_conf)
         with io.open(hba_conf, "r+") as conf:
             # get rid of comments and blank lines
             lines = [
@@ -585,10 +584,7 @@ def append_conf(self, line='', filename=PG_CONF_FILE, **kwargs):
             lines.append('{} = {}'.format(option, value))
 
         config_name = os.path.join(self.data_dir, filename)
-        with io.open(config_name, 'a') as conf:
-            for line in lines:
-                conf.write(text_type(line))
-                conf.write(text_type('\n'))
+        self.os_ops.write(config_name, '\n'.join(lines))
 
         return self
 

testgres/os_ops.py

@@ -0,0 +1,120 @@
+import os
+import shutil
+import subprocess
+
+import paramiko
+
+
+class OsOperations:
+
+    def __init__(self, host, ssh_key=None, username='dev'):
+        self.host = host
+        self.ssh_key = ssh_key
+        self.username = username
+        self.remote = self.host != 'localhost'
+        self.ssh = None
+
+        if self.remote:
+            self.ssh = self.connect()
+
+    def __del__(self):
+        if self.ssh:
+            self.ssh.close()
+
+    def connect(self):
+        # Check that the ssh key file exists and is a file
+        if not os.path.isfile(self.ssh_key):
+            raise ValueError(f"{self.ssh_key} does not exist or is not a file")
+
+        # Load the private key with the correct key type
+        with open(self.ssh_key, 'r') as f:
+            key_data = f.read()
+            if 'BEGIN OPENSSH PRIVATE KEY' in key_data:
+                key = paramiko.Ed25519Key.from_private_key_file(self.ssh_key)
+            else:
+                key = paramiko.RSAKey.from_private_key_file(self.ssh_key)
+
+        # Now use the loaded key in your SSH client code
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh.connect(self.host, username=self.username, pkey=key)
+        return ssh
+
+    def makedirs(self, path):
+        if self.remote:
+            # Remove the remote directory
+            stdin, stdout, stderr = self.ssh.exec_command(f'rm -rf {path}')
+            stdout.channel.recv_exit_status()
+
+            # Create the remote directory
+            stdin, stdout, stderr = self.ssh.exec_command(f'mkdir -p {path}')
+            stdout.channel.recv_exit_status()
+        else:
+            shutil.rmtree(path, ignore_errors=True)
+            os.makedirs(path)
+
+    def write(self, filename, text):
+        if self.remote:
+            command = f"echo '{text}' >> {filename}"
+            stdin, stdout, stderr = self.ssh.exec_command(command)
+            stdout.channel.recv_exit_status()
+        else:
+            with open(filename, 'a') as file:
+                file.write(text)
+
+    def read(self, filename):
+        if self.remote:
+            stdin, stdout, stderr = self.ssh.exec_command(f'cat {filename}')
+            stdout.channel.recv_exit_status()
+            output = stdout.read().decode().rstrip()
+        else:
+            with open(filename) as file:
+                output = file.read()
+        return output
+
+    def readlines(self, filename):
+        if self.remote:
+            stdin, stdout, stderr = self.ssh.exec_command(f'cat {filename}')
+            stdout.channel.recv_exit_status()
+            output = stdout.read().decode().rstrip()
+        else:
+            with open(filename) as file:
+                output = file.read()
+        return output
+
+    def get_name(self):
+        if self.remote:
+            stdin, stdout, stderr = self.ssh.exec_command('python -c "import os; print(os.name)"')
+            stdout.channel.recv_exit_status()
+            name = stdout.read().decode().strip()
+        else:
+            name = os.name
+        return name
+
+    def kill(self, pid, signal):
+        if self.remote:
+            stdin, stdout, stderr = self.ssh.exec_command(f'kill -{signal} {pid}')
+            stdout.channel.recv_exit_status()
+        else:
+            os.kill(pid, signal)
+
+    def environ(self, var_name):
+        if self.remote:
+            # Get env variable
+            check_command = f"echo ${var_name}"
+            stdin, stdout, stderr = self.ssh.exec_command(check_command)
+            stdout.channel.recv_exit_status()
+            var_val = stdout.read().strip().decode('utf-8')
+        else:
+            var_val = os.environ.get(var_name)
+        return var_val
+
+    def execute(self, cmd):
+        if self.remote:
+            stdin, stdout, stderr = self.ssh.exec_command(cmd)
+            stdout.channel.recv_exit_status()
+            result = stdout.read().decode('utf-8')
+        else:
+            result = subprocess.check_output([cmd]).decode('utf-8')
+        return result
+

testgres/utils.py

@@ -19,6 +19,7 @@
 # Add remote host call
 from fabric import Connection
 
+from os_ops import OsOperations
 from .config import testgres_config
 from .exceptions import ExecUtilException
 
@@ -51,7 +52,7 @@ def release_port(port):
     bound_ports.discard(port)
 
 
-def execute_utility(args, logfile=None, hostname='localhost', ssh_key=None):
+def execute_utility(args, logfile=None, host='localhost', ssh_key=None, user='dev'):
     """
     Execute utility wrapper (pg_ctl, pg_dump etc).
 
@@ -60,20 +61,19 @@ def execute_utility(args, logfile=None, hostname='localhost', ssh_key=None):
 
     Args:
         args: utility + arguments (list).
-        host_params : dict {
+        host: remote connection host
+        ssh_key: remote connection ssh_key
+        user: remote connection user
         logfile: path to file to store stdout and stderr.
 
     Returns:
         stdout of executed utility.
     """
-
-
-    if hostname != 'localhost':
+    if host != 'localhost':
         conn = Connection(
-            host=hostname,
-            user="dev",
+            host=host,
+            user=user,
             connect_kwargs={
-                # XXX pass SSH key path
                 "key_filename": ssh_key,
             },
         )
@@ -87,9 +87,8 @@ def execute_utility(args, logfile=None, hostname='localhost', ssh_key=None):
 
         return result
 
-    
     # run utility
-    if os.name == 'nt':
+    elif os.name == 'nt':
         # using output to a temporary file in Windows
         buf = tempfile.NamedTemporaryFile()
 
@@ -146,24 +145,25 @@ def execute_utility(args, logfile=None, hostname='localhost', ssh_key=None):
     return out
 
 
-def get_bin_path(filename):
+def get_bin_path(filename, host='localhost', ssh_key=None):
     """
     Return absolute path to an executable using PG_BIN or PG_CONFIG.
     This function does nothing if 'filename' is already absolute.
     """
+    os_ops = OsOperations(host, ssh_key)
 
     # check if it's already absolute
     if os.path.isabs(filename):
         return filename
 
     # try PG_CONFIG
-    pg_config = os.environ.get("PG_CONFIG")
+    pg_config = os_ops.environ("PG_CONFIG")
     if pg_config:
         bindir = get_pg_config()["BINDIR"]
         return os.path.join(bindir, filename)
 
     # try PG_BIN
-    pg_bin = os.environ.get("PG_BIN")
+    pg_bin = os_ops.environ("PG_BIN")
     if pg_bin:
         return os.path.join(pg_bin, filename)