Merge with master

asavchkov · asavchkov · commit 3f9c618f14d4 · 2024-06-26T23:40:56.000+07:00
diff --git a/testgres/node.py b/testgres/node.py
@@ -1641,7 +1641,7 @@ 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):
+    def upgrade_from(self, old_node, options=None):
         """
         Upgrade this node from an old node using pg_upgrade.
 
@@ -1654,6 +1654,9 @@ def upgrade_from(self, old_node):
         if not os.path.exists(self.data_dir):
             self.init()
 
+        if not options:
+            options = []
+
         pg_upgrade_binary = self._get_bin_path("pg_upgrade")
 
         if not os.path.exists(pg_upgrade_binary):
@@ -1668,6 +1671,7 @@ def upgrade_from(self, old_node):
             "--old-port", str(old_node.port),
             "--new-port", str(self.port),
         ]
+        upgrade_command += options
 
         return self.os_ops.exec_command(upgrade_command)
 
diff --git a/testgres/operations/remote_ops.py b/testgres/operations/remote_ops.py
@@ -48,13 +48,14 @@ def __init__(self, conn_params: ConnectionParams):
         super().__init__(conn_params.username)
         self.conn_params = conn_params
         self.host = conn_params.host
-        self.ssh_key = conn_params.ssh_key
         self.port = conn_params.port
+        self.ssh_key = conn_params.ssh_key
         self.ssh_cmd = ["-o StrictHostKeyChecking=no"]
+        self.ssh_args = []
         if self.ssh_key:
-            self.ssh_cmd += ["-i", self.ssh_key]
+            self.ssh_args += ["-i", self.ssh_key]
         if self.port:
-            self.ssh_cmd += ["-p", self.port]
+            self.ssh_args += ["-p", self.port]
         self.remote = True
         self.username = conn_params.username or self.get_user()
         self.tunnel_process = None
@@ -113,9 +114,9 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
         """
         ssh_cmd = []
         if isinstance(cmd, str):
-            ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + [cmd]
+            ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_args + [cmd]
         elif isinstance(cmd, list):
-            ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + cmd
+            ssh_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_args + cmd
         process = subprocess.Popen(ssh_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         if get_process:
             return process
@@ -264,9 +265,9 @@ def mkdtemp(self, prefix=None):
         - prefix (str): The prefix of the temporary directory name.
         """
         if prefix:
-            command = ["ssh" +  f"{self.username}@{self.host}"] + self.ssh_cmd + [f"mktemp -d {prefix}XXXXX"]
+            command = ["ssh"] + self.ssh_args + [f"{self.username}@{self.host}", f"mktemp -d {prefix}XXXXX"]
         else:
-            command = ["ssh", f"{self.username}@{self.host}"] + self.ssh_cmd + ["mktemp -d"]
+            command = ["ssh"] + self.ssh_args + [f"{self.username}@{self.host}", "mktemp -d"]
 
         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
 
@@ -309,11 +310,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
             mode = "r+b" if binary else "r+"
 
         with tempfile.NamedTemporaryFile(mode=mode, delete=False) as tmp_file:
-            # Because in scp we set up port using -P option
-            scp_ssh_cmd = ['-P' if x == '-p' else x for x in self.ssh_cmd]
+            # 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_ssh_cmd + [f"{self.username}@{self.host}:{filename}", tmp_file.name]
+                scp_cmd = ['scp'] + scp_args + [f"{self.username}@{self.host}:{filename}", tmp_file.name]
                 subprocess.run(scp_cmd, check=False)  # The file might not exist yet
                 tmp_file.seek(0, os.SEEK_END)
 
@@ -329,11 +330,11 @@ def write(self, filename, data, truncate=False, binary=False, read_and_write=Fal
                 tmp_file.write(data)
 
             tmp_file.flush()
-            scp_cmd = ['scp'] + scp_ssh_cmd + [tmp_file.name, f"{self.username}@{self.host}:{filename}"]
+            scp_cmd = ['scp'] + scp_args + [tmp_file.name, f"{self.username}@{self.host}:{filename}"]
             subprocess.run(scp_cmd, check=True)
 
             remote_directory = os.path.dirname(filename)
-            mkdir_cmd = ['ssh', f"{self.username}@{self.host}"] + self.ssh_cmd + [f"mkdir -p {remote_directory}"]
+            mkdir_cmd = ['ssh'] + self.ssh_args + [f"{self.username}@{self.host}", f"mkdir -p {remote_directory}"]
             subprocess.run(mkdir_cmd, check=True)
 
             os.remove(tmp_file.name)
@@ -398,7 +399,7 @@ def get_pid(self):
         return int(self.exec_command("echo $$", encoding=get_default_encoding()))
 
     def get_process_children(self, pid):
-        command = ["ssh", f"{self.username}@{self.host}"] + self.ssh_cmd + [f"pgrep -P {pid}"]
+        command = ["ssh"] + self.ssh_args + [f"{self.username}@{self.host}", f"pgrep -P {pid}"]
 
         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
 
diff --git a/testgres/plugins/pg_probackup2/pg_probackup2/app.py b/testgres/plugins/pg_probackup2/pg_probackup2/app.py
@@ -56,6 +56,7 @@ def __init__(self, test_class: unittest.TestCase,
         self.verbose = init_params.verbose
         self.archive_compress = init_params.archive_compress
         self.test_class.output = None
+        self.execution_time = None
 
     def run(self, command, gdb=False, old_binary=False, return_id=True, env=None,
             skip_log_directory=False, expect_error=False, use_backup_dir=True):
@@ -113,16 +114,17 @@ def run(self, command, gdb=False, old_binary=False, return_id=True, env=None,
                 cmdline = ['gdbserver'] + ['localhost:' + str(gdb_port)] + cmdline
                 print("pg_probackup gdb suspended, waiting gdb connection on localhost:{0}".format(gdb_port))
 
+            start_time = time.time()
             self.test_class.output = subprocess.check_output(
                 cmdline,
                 stderr=subprocess.STDOUT,
                 env=env
             ).decode('utf-8', errors='replace')
+            end_time = time.time()
+            self.execution_time = end_time - start_time
+
             if command[0] == 'backup' and return_id:
-                # return backup ID
-                for line in self.test_class.output.splitlines():
-                    if 'INFO: Backup' and 'completed' in line:
-                        result = line.split()[2]
+                result = self.get_backup_id()
             else:
                 result = self.test_class.output
             if expect_error is True:
@@ -139,6 +141,19 @@ def run(self, command, gdb=False, old_binary=False, return_id=True, env=None,
             else:
                 raise ProbackupException(self.test_class.output, self.test_class.cmd)
 
+    def get_backup_id(self):
+        if init_params.major_version > 2:
+            pattern = re.compile(r"Backup (.*) completed successfully.")
+            for line in self.test_class.output.splitlines():
+                match = pattern.search(line)
+                if match:
+                    return match.group(1)
+        else:
+            for line in self.test_class.output.splitlines():
+                if 'INFO: Backup' and 'completed' in line:
+                    return line.split()[2]
+        return None
+
     def init(self, options=None, old_binary=False, skip_log_directory=False, expect_error=False, use_backup_dir=True):
         if options is None:
             options = []
diff --git a/testgres/plugins/pg_probackup2/setup.py b/testgres/plugins/pg_probackup2/setup.py
@@ -4,7 +4,7 @@
     from distutils.core import setup
 
 setup(
-    version='0.0.2',
+    version='0.0.3',
     name='testgres_pg_probackup2',
     packages=['pg_probackup2', 'pg_probackup2.storage'],
     description='Plugin for testgres that manages pg_probackup2',
