PBCKP-152 change local function on execution by ssh

Author: v.shepard

setup.py

@@ -6,7 +6,7 @@
     from distutils.core import setup
 
 # Basic dependencies
-install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil"]
+install_requires = ["pg8000", "port-for>=0.4", "six>=1.9.0", "psutil", "fabric"]
 
 # Add compatibility enum class
 if sys.version_info < (3, 4):
@@ -21,7 +21,7 @@
     readme = f.read()
 
 setup(
-    version='1.8.6',
+    version='1.9.0',
     name='testgres',
     packages=['testgres'],
     description='Testing utility for PostgreSQL and its extensions',

testgres/api.py

@@ -30,13 +30,19 @@
 PostgresNode(name='...', port=..., base_dir='...')
 [(3,)]
 """
+from defaults import default_username
 from .node import PostgresNode
 
 
 def get_new_node(name=None, base_dir=None, **kwargs):
     """
     Simply a wrapper around :class:`.PostgresNode` constructor.
     See :meth:`.PostgresNode.__init__` for details.
+    For remote connection you can add next parameters:
+    host='127.0.0.1',
+    hostname='localhost',
+    ssh_key=None,
+    username=default_username()
     """
     # NOTE: leave explicit 'name' and 'base_dir' for compatibility
     return PostgresNode(name=name, base_dir=base_dir, **kwargs)

testgres/backup.py

@@ -6,6 +6,7 @@
 from six import raise_from
 from tempfile import mkdtemp
 
+from os_ops import OsOperations
 from .enums import XLogMethod
 
 from .consts import \
@@ -47,7 +48,7 @@ def __init__(self,
             username: database user name.
             xlog_method: none | fetch | stream (see docs)
         """
-
+        self.os_ops = node.os_ops
         if not node.status():
             raise BackupException('Node must be running')
 
@@ -60,8 +61,8 @@ def __init__(self,
                 raise BackupException(msg)
 
         # Set default arguments
-        username = username or default_username()
-        base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP)
+        username = username or self.os_ops.get_user()
+        base_dir = base_dir or self.os_ops.mkdtemp(prefix=TMP_BACKUP)
 
         # public
         self.original_node = node
@@ -81,7 +82,7 @@ def __init__(self,
             "-D", data_dir,
             "-X", xlog_method.value
         ]  # yapf: disable
-        execute_utility(_params, self.log_file)
+        execute_utility(_params, self.log_file, hostname=node.hostname, ssh_key=node.ssh_key)
 
     def __enter__(self):
         return self
@@ -107,7 +108,7 @@ def _prepare_dir(self, destroy):
         available = not destroy
 
         if available:
-            dest_base_dir = mkdtemp(prefix=TMP_NODE)
+            dest_base_dir = self.os_ops.mkdtemp(prefix=TMP_NODE)
 
             data1 = os.path.join(self.base_dir, DATA_DIR)
             data2 = os.path.join(dest_base_dir, DATA_DIR)
@@ -185,4 +186,4 @@ def cleanup(self):
 
         if self._available:
             self._available = False
-            rmtree(self.base_dir, ignore_errors=True)
+            self.os_ops.rmdirs(self.base_dir, ignore_errors=True)

testgres/cache.py

@@ -6,6 +6,7 @@
 from shutil import copytree
 from six import raise_from
 
+from os_ops import OsOperations
 from .config import testgres_config
 
 from .consts import XLOG_CONTROL_FILE
@@ -21,14 +22,16 @@
     execute_utility
 
 
-def cached_initdb(data_dir, logfile=None, params=None):
+def cached_initdb(data_dir, logfile=None, hostname='localhost', ssh_key=None, params=None):
     """
     Perform initdb or use cached node files.
     """
+    os_ops = OsOperations(hostname=hostname, ssh_key=ssh_key)
+
     def call_initdb(initdb_dir, log=None):
         try:
             _params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
-            execute_utility(_params + (params or []), log)
+            execute_utility(_params + (params or []), log, hostname=hostname, ssh_key=ssh_key)
         except ExecUtilException as e:
             raise_from(InitNodeException("Failed to run initdb"), e)
 
@@ -39,26 +42,27 @@ def call_initdb(initdb_dir, log=None):
         cached_data_dir = testgres_config.cached_initdb_dir
 
         # Initialize cached initdb
-        if not os.path.exists(cached_data_dir) or \
-           not os.listdir(cached_data_dir):
+
+        if not os_ops.path_exists(cached_data_dir) or \
+           not os_ops.listdir(cached_data_dir):
             call_initdb(cached_data_dir)
 
         try:
             # Copy cached initdb to current data dir
-            copytree(cached_data_dir, data_dir)
+            os_ops.copytree(cached_data_dir, data_dir)
 
             # Assign this node a unique system id if asked to
             if testgres_config.cached_initdb_unique:
                 # XXX: write new unique system id to control file
                 # Some users might rely upon unique system ids, but
                 # our initdb caching mechanism breaks this contract.
                 pg_control = os.path.join(data_dir, XLOG_CONTROL_FILE)
-                with io.open(pg_control, "r+b") as f:
-                    f.write(generate_system_id())    # overwrite id
+                system_id = generate_system_id()
+                os_ops.write(pg_control, system_id, truncate=True, binary=True, read_and_write=True)
 
                 # XXX: build new WAL segment with our system id
                 _params = [get_bin_path("pg_resetwal"), "-D", data_dir, "-f"]
-                execute_utility(_params, logfile)
+                execute_utility(_params, logfile, hostname=hostname, ssh_key=ssh_key)
 
         except ExecUtilException as e:
             msg = "Failed to reset WAL for system id"

testgres/config.py

@@ -43,6 +43,9 @@ class GlobalConfig(object):
 
     _cached_initdb_dir = None
     """ underlying class attribute for cached_initdb_dir property """
+
+    os_ops = None
+    """ OsOperation object that allows work on remote host """
     @property
     def cached_initdb_dir(self):
         """ path to a temp directory for cached initdb. """
@@ -52,8 +55,15 @@ def cached_initdb_dir(self):
     def cached_initdb_dir(self, value):
         self._cached_initdb_dir = value
 
+        # NOTE: assign initial cached dir for initdb
+        if self.os_ops:
+            testgres_config.cached_initdb_dir = self.os_ops.mkdtemp(prefix=TMP_CACHE)
+        else:
+            testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
+
         if value:
             cached_initdb_dirs.add(value)
+        return testgres_config.cached_initdb_dir
 
     @property
     def temp_dir(self):
@@ -133,9 +143,12 @@ def copy(self):
 
 
 @atexit.register
-def _rm_cached_initdb_dirs():
+def _rm_cached_initdb_dirs(os_ops=None):
     for d in cached_initdb_dirs:
-        rmtree(d, ignore_errors=True)
+        if os_ops:
+            os_ops.rmtree(d, ignore_errors=True)
+        else:
+            rmtree(d, ignore_errors=True)
 
 
 def push_config(**options):
@@ -195,7 +208,3 @@ def configure_testgres(**options):
     """
 
     testgres_config.update(options)
-
-
-# NOTE: assign initial cached dir for initdb
-testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)

testgres/connection.py

@@ -1,4 +1,5 @@
 # coding: utf-8
+from os_ops import OsOperations
 
 # we support both pg8000 and psycopg2
 try:
@@ -41,11 +42,12 @@ def __init__(self,
 
         self._node = node
 
-        self._connection = pglib.connect(database=dbname,
-                                         user=username,
-                                         password=password,
-                                         host=node.host,
-                                         port=node.port)
+        self.os_ops = OsOperations(node.host, node.hostname, node.ssh_key, node.username)
+        self._connection = self.os_ops.db_connect(dbname=dbname,
+                                                  user=username,
+                                                  password=password,
+                                                  host=node.host,
+                                                  port=node.port)
 
         self._connection.autocommit = autocommit
         self._cursor = self.connection.cursor()
@@ -102,17 +104,24 @@ def rollback(self):
         return self
 
     def execute(self, query, *args):
-        self.cursor.execute(query, args)
-
         try:
-            res = self.cursor.fetchall()
-
-            # pg8000 might return tuples
-            if isinstance(res, tuple):
-                res = [tuple(t) for t in res]
-
-            return res
-        except Exception:
+            with self.connection.cursor() as cursor:
+                cursor.execute(query, args)
+                try:
+                    res = cursor.fetchall()
+
+                    # pg8000 might return tuples
+                    if isinstance(res, tuple):
+                        res = [tuple(t) for t in res]
+
+                    return res
+                except (pglib.ProgrammingError, pglib.InternalError) as e:
+                    # An error occurred while trying to fetch results (e.g., no results to fetch)
+                    print(f"Error fetching results: {e}")
+                    return None
+        except (pglib.Error, Exception) as e:
+            # Handle other database errors
+            print(f"Error executing query: {e}")
             return None
 
     def close(self):

testgres/defaults.py

@@ -13,12 +13,15 @@ def default_dbname():
     return 'postgres'
 
 
-def default_username():
+def default_username(os_ops=None):
     """
     Return default username (current user).
     """
-
-    return getpass.getuser()
+    if os_ops:
+        user = os_ops.get_user()
+    else:
+        user = getpass.getuser()
+    return user
 
 
 def generate_app_name():
@@ -29,7 +32,7 @@ def generate_app_name():
     return 'testgres-{}'.format(str(uuid.uuid4()))
 
 
-def generate_system_id():
+def generate_system_id(os_ops=None):
     """
     Generate a new 64-bit unique system identifier for node.
     """
@@ -44,7 +47,10 @@ def generate_system_id():
     system_id = 0
     system_id |= (secs << 32)
     system_id |= (usecs << 12)
-    system_id |= (os.getpid() & 0xFFF)
+    if os_ops:
+        system_id |= (os_ops.get_pid() & 0xFFF)
+    else:
+        system_id |= (os.getpid() & 0xFFF)
 
     # pack ULL in native byte order
     return struct.pack('=Q', system_id)

testgres/logger.py

@@ -6,6 +6,20 @@
 import time
 
 
+# create logger
+log = logging.getLogger('Testgres')
+log.setLevel(logging.DEBUG)
+# create console handler and set level to debug
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+# create formatter
+formatter = logging.Formatter('\n%(asctime)s - %(name)s[%(levelname)s]: %(message)s')
+# add formatter to ch
+ch.setFormatter(formatter)
+# add ch to logger
+log.addHandler(ch)
+
+
 class TestgresLogger(threading.Thread):
     """
     Helper class to implement reading from log files.

testgres/node.py

@@ -12,6 +12,7 @@
 import subprocess
 import time
 
+from os_ops import OsOperations
 
 try:
     from collections.abc import Iterable
@@ -129,7 +130,8 @@ def __repr__(self):
 
 
 class PostgresNode(object):
-    def __init__(self, name=None, port=None, base_dir=None):
+    def __init__(self, name=None, port=None, base_dir=None,
+                 host='127.0.0.1', hostname='localhost', ssh_key=None, username=default_username()):
         """
         PostgresNode constructor.
 
@@ -147,10 +149,14 @@ def __init__(self, name=None, port=None, base_dir=None):
         self._master = None
 
         # basic
-        self.host = '127.0.0.1'
         self.name = name or generate_app_name()
         self.port = port or reserve_port()
 
+        self.host = host
+        self.hostname = hostname
+        self.ssh_key = ssh_key
+        self.os_ops = OsOperations(host, hostname, ssh_key, username=username)
+
         # defaults for __exit__()
         self.cleanup_on_good_exit = testgres_config.node_cleanup_on_good_exit
         self.cleanup_on_bad_exit = testgres_config.node_cleanup_on_bad_exit
@@ -455,9 +461,12 @@ def init(self, initdb_params=None, **kwargs):
         """
 
         # initialize this PostgreSQL node
-        cached_initdb(data_dir=self.data_dir,
-                      logfile=self.utils_log_file,
-                      params=initdb_params)
+        cached_initdb(
+            data_dir=self.data_dir,
+            logfile=self.utils_log_file,
+            hostname=self.hostname,
+            ssh_key=self.ssh_key,
+            params=initdb_params)
 
         # initialize default config files
         self.default_conf(**kwargs)
@@ -514,6 +523,10 @@ def get_auth_method(t):
                 new_lines = [
                     u"local\treplication\tall\t\t\t{}\n".format(auth_local),
                     u"host\treplication\tall\t127.0.0.1/32\t{}\n".format(auth_host),
+
+                    u"host\treplication\tall\t0.0.0.0/0\t{}\n".format(auth_host),
+                    u"host\tall\tall\t0.0.0.0/0\t{}\n".format(auth_host),
+
                     u"host\treplication\tall\t::1/128\t\t{}\n".format(auth_host)
                 ]  # yapf: disable