introduce clean_on_error() for safer node handling

Author: funbringer

testgres/backup.py

@@ -21,7 +21,8 @@
 
 from .utils import \
     get_bin_path, \
-    execute_utility
+    execute_utility, \
+    clean_on_error
 
 
 class NodeBackup(object):
@@ -56,8 +57,8 @@ def __init__(self,
             try:
                 xlog_method = XLogMethod(xlog_method)
             except ValueError:
-                raise BackupException(
-                    'Invalid xlog_method "{}"'.format(xlog_method))
+                msg = 'Invalid xlog_method "{}"'.format(xlog_method)
+                raise BackupException(msg)
 
         # Set default arguments
         username = username or default_username()
@@ -144,15 +145,16 @@ def spawn_primary(self, name=None, destroy=True):
 
         # Build a new PostgresNode
         from .node import PostgresNode
-        node = PostgresNode(name=name, base_dir=base_dir)
+        with clean_on_error(PostgresNode(name=name,
+                                         base_dir=base_dir)) as node:
 
-        # New nodes should always remove dir tree
-        node._should_rm_dirs = True
+            # New nodes should always remove dir tree
+            node._should_rm_dirs = True
 
-        node.append_conf(PG_CONF_FILE, "\n")
-        node.append_conf(PG_CONF_FILE, "port = {}".format(node.port))
+            node.append_conf(PG_CONF_FILE, "\n")
+            node.append_conf(PG_CONF_FILE, "port = {}".format(node.port))
 
-        return node
+            return node
 
     def spawn_replica(self, name=None, destroy=True, slot=None):
         """
@@ -168,13 +170,14 @@ def spawn_replica(self, name=None, destroy=True, slot=None):
         """
 
         # Build a new PostgresNode
-        node = self.spawn_primary(name=name, destroy=destroy)
+        with clean_on_error(self.spawn_primary(name=name,
+                                               destroy=destroy)) as node:
 
-        # Assign it a master and a recovery file (private magic)
-        node._assign_master(self.original_node)
-        node._create_recovery_conf(username=self.username, slot=slot)
+            # Assign it a master and a recovery file (private magic)
+            node._assign_master(self.original_node)
+            node._create_recovery_conf(username=self.username, slot=slot)
 
-        return node
+            return node
 
     def cleanup(self):
         if self._available:

testgres/node.py

@@ -61,7 +61,8 @@
     pg_version_ge, \
     reserve_port, \
     release_port, \
-    execute_utility
+    execute_utility, \
+    clean_on_error
 
 from .backup import NodeBackup
 
@@ -979,10 +980,9 @@ def replicate(self, name=None, slot=None, **kwargs):
             base_dir: the base directory for data files and logs
         """
 
-        backup = self.backup(**kwargs)
-
         # transform backup into a replica
-        return backup.spawn_replica(name=name, destroy=True, slot=slot)
+        with clean_on_error(self.backup(**kwargs)) as backup:
+            return backup.spawn_replica(name=name, destroy=True, slot=slot)
 
     def catchup(self, dbname=None, username=None):
         """

testgres/utils.py

@@ -9,6 +9,7 @@
 import subprocess
 import sys
 
+from contextlib import contextmanager
 from distutils.version import LooseVersion
 
 from .config import testgres_config
@@ -217,4 +218,23 @@ def file_tail(f, num_lines):
 
 
 def eprint(*args, **kwargs):
+    """
+    Print stuff to stderr.
+    """
+
     print(*args, file=sys.stderr, **kwargs)
+
+
+@contextmanager
+def clean_on_error(node):
+    """
+    Context manager to wrap PostgresNode and such.
+    Calls cleanup() method when underlying code raises an exception.
+    """
+
+    try:
+        yield node
+    except Exception:
+        # TODO: should we wrap this in try-block?
+        node.cleanup()
+        raise