Add function pg_update

vshepard · vshepard · commit 334bb226d690 · 2024-01-02T08:05:30.000+01:00
diff --git a/testgres/cache.py b/testgres/cache.py
@@ -22,19 +22,20 @@
 from .operations.os_ops import OsOperations
 
 
-def cached_initdb(data_dir, logfile=None, params=None, os_ops: OsOperations = LocalOperations()):
+def cached_initdb(data_dir, logfile=None, params=None, os_ops: OsOperations = LocalOperations(), bin_path=None, cached=True):
     """
     Perform initdb or use cached node files.
     """
 
     def call_initdb(initdb_dir, log=logfile):
         try:
-            _params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
+            initdb_path = os.path.join(bin_path, 'initdb') if bin_path else get_bin_path("initdb")
+            _params = [initdb_path, "-D", initdb_dir, "-N"]
             execute_utility(_params + (params or []), log)
         except ExecUtilException as e:
             raise_from(InitNodeException("Failed to run initdb"), e)
 
-    if params or not testgres_config.cache_initdb:
+    if params or not testgres_config.cache_initdb or not cached:
         call_initdb(data_dir, logfile)
     else:
         # Fetch cached initdb dir
diff --git a/testgres/node.py b/testgres/node.py
@@ -127,20 +127,23 @@ def __repr__(self):
 
 
 class PostgresNode(object):
-    def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionParams = ConnectionParams()):
+    def __init__(self, name=None, port=None, base_dir=None, conn_params: ConnectionParams = ConnectionParams(), bin_dir=None, prefix=None):
         """
         PostgresNode constructor.
 
         Args:
             name: node's application name.
             port: port to accept connections.
             base_dir: path to node's data directory.
+            bin_dir: path to node's binary directory.
         """
 
         # private
-        self._pg_version = PgVer(get_pg_version())
+        self._pg_version = PgVer(get_pg_version(bin_dir))
         self._should_free_port = port is None
         self._base_dir = base_dir
+        self._bin_dir = bin_dir
+        self._prefix = prefix
         self._logger = None
         self._master = None
 
@@ -281,14 +284,20 @@ def master(self):
     @property
     def base_dir(self):
         if not self._base_dir:
-            self._base_dir = self.os_ops.mkdtemp(prefix=TMP_NODE)
+            self._base_dir = self.os_ops.mkdtemp(prefix=self._prefix or TMP_NODE)
 
         # NOTE: it's safe to create a new dir
         if not self.os_ops.path_exists(self._base_dir):
             self.os_ops.makedirs(self._base_dir)
 
         return self._base_dir
 
+    @property
+    def bin_dir(self):
+        if not self._bin_dir:
+            self._bin_dir = os.path.dirname(get_bin_path("pg_config"))
+        return self._bin_dir
+
     @property
     def logs_dir(self):
         path = os.path.join(self.base_dir, LOGS_DIR)
@@ -441,7 +450,7 @@ def _collect_special_files(self):
 
         return result
 
-    def init(self, initdb_params=None, **kwargs):
+    def init(self, initdb_params=None, cached=True, **kwargs):
         """
         Perform initdb for this node.
 
@@ -460,7 +469,9 @@ def init(self, initdb_params=None, **kwargs):
             data_dir=self.data_dir,
             logfile=self.utils_log_file,
             os_ops=self.os_ops,
-            params=initdb_params)
+            params=initdb_params,
+            bin_path=self.bin_dir,
+            cached=False)
 
         # initialize default config files
         self.default_conf(**kwargs)
@@ -619,7 +630,7 @@ def status(self):
 
         try:
             _params = [
-                get_bin_path("pg_ctl"),
+                self._get_bin_path('pg_ctl'),
                 "-D", self.data_dir,
                 "status"
             ]  # yapf: disable
@@ -645,7 +656,7 @@ def get_control_data(self):
         """
 
         # this one is tricky (blame PG 9.4)
-        _params = [get_bin_path("pg_controldata")]
+        _params = [self._get_bin_path("pg_controldata")]
         _params += ["-D"] if self._pg_version >= PgVer('9.5') else []
         _params += [self.data_dir]
 
@@ -708,7 +719,7 @@ def start(self, params=[], wait=True):
             return self
 
         _params = [
-            get_bin_path("pg_ctl"),
+            self._get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "-l", self.pg_log_file,
             "-w" if wait else '-W',  # --wait or --no-wait
@@ -742,7 +753,7 @@ def stop(self, params=[], wait=True):
             return self
 
         _params = [
-            get_bin_path("pg_ctl"),
+            self._get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "-w" if wait else '-W',  # --wait or --no-wait
             "stop"
@@ -782,7 +793,7 @@ def restart(self, params=[]):
         """
 
         _params = [
-            get_bin_path("pg_ctl"),
+            self._get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "-l", self.pg_log_file,
             "-w",  # wait
@@ -814,7 +825,7 @@ def reload(self, params=[]):
         """
 
         _params = [
-            get_bin_path("pg_ctl"),
+            self._get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "reload"
         ] + params  # yapf: disable
@@ -835,7 +846,7 @@ def promote(self, dbname=None, username=None):
         """
 
         _params = [
-            get_bin_path("pg_ctl"),
+            self._get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "-w",  # wait
             "promote"
@@ -871,7 +882,7 @@ def pg_ctl(self, params):
         """
 
         _params = [
-            get_bin_path("pg_ctl"),
+            self._get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "-w"  # wait
         ] + params  # yapf: disable
@@ -945,7 +956,7 @@ def psql(self,
         username = username or default_username()
 
         psql_params = [
-            get_bin_path("psql"),
+            self._get_bin_path("psql"),
             "-p", str(self.port),
             "-h", self.host,
             "-U", username,
@@ -1066,7 +1077,7 @@ def tmpfile():
         filename = filename or tmpfile()
 
         _params = [
-            get_bin_path("pg_dump"),
+            self._get_bin_path("pg_dump"),
             "-p", str(self.port),
             "-h", self.host,
             "-f", filename,
@@ -1094,7 +1105,7 @@ def restore(self, filename, dbname=None, username=None):
         username = username or default_username()
 
         _params = [
-            get_bin_path("pg_restore"),
+            self._get_bin_path("pg_restore"),
             "-p", str(self.port),
             "-h", self.host,
             "-U", username,
@@ -1364,7 +1375,7 @@ def pgbench(self,
         username = username or default_username()
 
         _params = [
-            get_bin_path("pgbench"),
+            self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
             "-U", username,
@@ -1416,7 +1427,7 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
         username = username or default_username()
 
         _params = [
-            get_bin_path("pgbench"),
+            self._get_bin_path("pgbench"),
             "-p", str(self.port),
             "-h", self.host,
             "-U", username,
@@ -1587,6 +1598,43 @@ def set_auto_conf(self, options, config='postgresql.auto.conf', rm_options={}):
 
         self.os_ops.write(path, auto_conf, truncate=True)
 
+    def upgrade_from(self, old_node):
+        """
+        Upgrade this node from an old node using pg_upgrade.
+
+        Args:
+            old_node: An instance of PostgresNode representing the old node.
+        """
+        if not os.path.exists(old_node.data_dir):
+            raise Exception("Old node must be initialized")
+
+        if not os.path.exists(self.data_dir):
+            self.init()
+
+        pg_upgrade_binary = self._get_bin_path("pg_upgrade")
+
+        if not os.path.exists(pg_upgrade_binary):
+            raise Exception("pg_upgrade does not exist in the new node's binary path")
+
+        upgrade_command = [
+            pg_upgrade_binary,
+            "--old-bindir", old_node.bin_dir,
+            "--new-bindir", self.bin_dir,
+            "--old-datadir", old_node.data_dir,
+            "--new-datadir", self.data_dir,
+            "--old-port", str(old_node.port),
+            "--new-port", str(self.port),
+        ]
+
+        return self.os_ops.exec_command(upgrade_command)
+
+    def _get_bin_path(self, filename):
+        if self.bin_dir:
+            bin_path = os.path.join(self.bin_dir, filename)
+        else:
+            bin_path = get_bin_path(filename)
+        return bin_path
+
 
 class NodeApp:
 
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -44,7 +44,7 @@ def __init__(self, conn_params=None):
     def _raise_exec_exception(message, command, exit_code, output):
         """Raise an ExecUtilException."""
         raise ExecUtilException(message=message.format(output),
-                                command=command,
+                                command=' '.join(command) if isinstance(command, list) else command,
                                 exit_code=exit_code,
                                 out=output)
 
diff --git a/testgres/utils.py b/testgres/utils.py
@@ -172,13 +172,14 @@ def cache_pg_config_data(cmd):
     return cache_pg_config_data("pg_config")
 
 
-def get_pg_version():
+def get_pg_version(bin_dir=None):
     """
     Return PostgreSQL version provided by postmaster.
     """
 
     # get raw version (e.g. postgres (PostgreSQL) 9.5.7)
-    _params = [get_bin_path('postgres'), '--version']
+    postgres_path = os.path.join(bin_dir, 'postgres') if bin_dir else get_bin_path('postgres')
+    _params = [postgres_path, '--version']
     raw_ver = tconf.os_ops.exec_command(_params, encoding='utf-8')
 
     # Remove "(Homebrew)" if present
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -1010,6 +1010,19 @@ def test_child_process_dies(self):
             # try to handle children list -- missing processes will have ptype "ProcessType.Unknown"
             [ProcessProxy(p) for p in children]
 
+    def test_upgrade_node(self):
+        old_bin_dir = os.path.dirname(get_bin_path("pg_config"))
+        new_bin_dir = os.path.dirname(get_bin_path("pg_config"))
+        node_old = get_new_node(prefix='node_old', bin_dir=old_bin_dir)
+        node_old.init()
+        node_old.start()
+        node_old.stop()
+        node_new = get_new_node(prefix='node_new', bin_dir=new_bin_dir)
+        node_new.init(cached=False)
+        res = node_new.upgrade_from(old_node=node_old)
+        node_new.start()
+        self.assertTrue(b'Upgrade Complete' in res)
+
 
 if __name__ == '__main__':
     if os.environ.get('ALT_CONFIG'):
