Fix node cleanup (#135)

demonolock · web-flow · commit 1d68e91a71f9 · 2024-07-01T12:35:39.000+02:00
* Add expect_error to pg_upgrade

* Fix node cleanup - rmdirs

* Add cleanup parameter - clean full dir
diff --git a/testgres/node.py b/testgres/node.py
@@ -914,13 +914,14 @@ def free_port(self):
             self._should_free_port = False
             release_port(self.port)
 
-    def cleanup(self, max_attempts=3):
+    def cleanup(self, max_attempts=3, full=False):
         """
         Stop node if needed and remove its data/logs directory.
         NOTE: take a look at TestgresConfig.node_cleanup_full.
 
         Args:
             max_attempts: how many times should we try to stop()?
+            full: clean full base dir
 
         Returns:
             This instance of :class:`.PostgresNode`.
@@ -929,12 +930,12 @@ def cleanup(self, max_attempts=3):
         self._try_shutdown(max_attempts)
 
         # choose directory to be removed
-        if testgres_config.node_cleanup_full:
+        if testgres_config.node_cleanup_full or full:
             rm_dir = self.base_dir    # everything
         else:
             rm_dir = self.data_dir    # just data, save logs
 
-        self.os_ops.rmdirs(rm_dir, ignore_errors=True)
+        self.os_ops.rmdirs(rm_dir, ignore_errors=False)
 
         return self
 
@@ -1629,7 +1630,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, options=None):
+    def upgrade_from(self, old_node, options=None, expect_error=False):
         """
         Upgrade this node from an old node using pg_upgrade.
 
@@ -1657,11 +1658,11 @@ def upgrade_from(self, old_node, options=None):
             "--old-datadir", old_node.data_dir,
             "--new-datadir", self.data_dir,
             "--old-port", str(old_node.port),
-            "--new-port", str(self.port),
+            "--new-port", str(self.port)
         ]
         upgrade_command += options
 
-        return self.os_ops.exec_command(upgrade_command)
+        return self.os_ops.exec_command(upgrade_command, expect_error=expect_error)
 
     def _get_bin_path(self, filename):
         if self.bin_dir:
diff --git a/testgres/operations/local_ops.py b/testgres/operations/local_ops.py
@@ -4,6 +4,7 @@
 import stat
 import subprocess
 import tempfile
+import time
 
 import psutil
 
@@ -19,13 +20,18 @@
 
 CMD_TIMEOUT_SEC = 60
 error_markers = [b'error', b'Permission denied', b'fatal']
+err_out_markers = [b'Failure']
 
 
-def has_errors(output):
+def has_errors(output=None, error=None):
     if output:
         if isinstance(output, str):
             output = output.encode(get_default_encoding())
-        return any(marker in output for marker in error_markers)
+        return any(marker in output for marker in err_out_markers)
+    if error:
+        if isinstance(error, str):
+            error = error.encode(get_default_encoding())
+        return any(marker in error for marker in error_markers)
     return False
 
 
@@ -107,8 +113,8 @@ def exec_command(self, cmd, wait_exit=False, verbose=False, expect_error=False,
         process, output, error = self._run_command(cmd, shell, input, stdin, stdout, stderr, get_process, timeout, encoding)
         if get_process:
             return process
-        if process.returncode != 0 or (has_errors(error) and not expect_error):
-            self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode, error)
+        if (process.returncode != 0 or has_errors(output=output, error=error)) and not expect_error:
+            self._raise_exec_exception('Utility exited with non-zero code. Error `{}`', cmd, process.returncode, error or output)
 
         if verbose:
             return process.returncode, output, error
@@ -142,8 +148,27 @@ def makedirs(self, path, remove_existing=False):
         except FileExistsError:
             pass
 
-    def rmdirs(self, path, ignore_errors=True):
-        return rmtree(path, ignore_errors=ignore_errors)
+    def rmdirs(self, path, ignore_errors=True, retries=3, delay=1):
+        """
+        Removes a directory and its contents, retrying on failure.
+
+        :param path: Path to the directory.
+        :param ignore_errors: If True, ignore errors.
+        :param retries: Number of attempts to remove the directory.
+        :param delay: Delay between attempts in seconds.
+        """
+        for attempt in range(retries):
+            try:
+                rmtree(path, ignore_errors=ignore_errors)
+                if not os.path.exists(path):
+                    return True
+            except FileNotFoundError:
+                return True
+            except Exception as e:
+                print(f"Error: Failed to remove directory {path} on attempt {attempt + 1}: {e}")
+            time.sleep(delay)
+        print(f"Error: Failed to remove directory {path} after {retries} attempts.")
+        return False
 
     def listdir(self, path):
         return os.listdir(path)
