Initial implementation of config stack

testgres/__init__.py

@@ -1,6 +1,12 @@
 from .api import get_new_node
 from .backup import NodeBackup
-from .config import TestgresConfig, configure_testgres
+
+from .config import \
+    TestgresConfig, \
+    configure_testgres, \
+    scoped_config, \
+    push_config, \
+    pop_config
 
 from .connection import \
     IsolationLevel, \

testgres/cache.py

@@ -1,13 +1,11 @@
 # coding: utf-8
 
-import atexit
 import os
 import shutil
-import tempfile
 
 from six import raise_from
 
-from .config import TestgresConfig
+from .config import testgres_config
 
 from .exceptions import \
     InitNodeException, \
@@ -18,44 +16,31 @@
     execute_utility
 
 
-def cached_initdb(data_dir, initdb_logfile, initdb_params=[]):
+def cached_initdb(data_dir, logfile=None, params=None):
     """
     Perform initdb or use cached node files.
     """
 
-    def call_initdb(initdb_dir):
+    def call_initdb(initdb_dir, log=None):
         try:
             _params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
-            execute_utility(_params + initdb_params, initdb_logfile)
+            execute_utility(_params + (params or []), log)
         except ExecUtilException as e:
             raise_from(InitNodeException("Failed to run initdb"), e)
 
-    def rm_cached_data_dir(cached_data_dir):
-        shutil.rmtree(cached_data_dir, ignore_errors=True)
-
-    # Call initdb if we have custom params or shouldn't cache it
-    if initdb_params or not TestgresConfig.cache_initdb:
-        call_initdb(data_dir)
+    if params or not testgres_config.cache_initdb:
+        call_initdb(data_dir, logfile)
     else:
-        # Set default temp dir for cached initdb
-        if TestgresConfig.cached_initdb_dir is None:
-
-            # Create default temp dir
-            TestgresConfig.cached_initdb_dir = tempfile.mkdtemp()
-
-            # Schedule cleanup
-            atexit.register(rm_cached_data_dir,
-                            TestgresConfig.cached_initdb_dir)
-
         # Fetch cached initdb dir
-        cached_data_dir = TestgresConfig.cached_initdb_dir
+        cached_data_dir = testgres_config.cached_initdb_dir
 
         # Initialize cached initdb
-        if not os.listdir(cached_data_dir):
+        if not os.path.exists(cached_data_dir) or \
+           not os.listdir(cached_data_dir):
             call_initdb(cached_data_dir)
 
         try:
             # Copy cached initdb to current data dir
             shutil.copytree(cached_data_dir, data_dir)
         except Exception as e:
-            raise_from(InitNodeException("Failed to copy files"), e)
+            raise_from(InitNodeException("Failed to spawn a node"), e)

testgres/config.py

@@ -1,26 +1,35 @@
 # coding: utf-8
 
+import atexit
+import copy
+import shutil
+import tempfile
 
-class TestgresConfig:
+from contextlib import contextmanager
+
+
+class GlobalConfig(object):
     """
     Global config (override default settings).
 
     Attributes:
-        cache_initdb:       shall we use cached initdb instance?
-        cached_initdb_dir:  shall we create a temp dir for cached initdb?
+        cache_initdb:           shall we use cached initdb instance?
+        cached_initdb_dir:      shall we create a temp dir for cached initdb?
 
-        cache_pg_config:    shall we cache pg_config results?
+        cache_pg_config:        shall we cache pg_config results?
 
-        use_python_logging: use python logging configuration for all nodes.
-        error_log_lines:    N of log lines to be included into exception (0=inf).
+        use_python_logging:     use python logging configuration for all nodes.
+        error_log_lines:        N of log lines to be shown in exception (0=inf).
 
-        node_cleanup_full:  shall we remove EVERYTHING (including logs)?
+        node_cleanup_full:          shall we remove EVERYTHING (including logs)?
         node_cleanup_on_good_exit:  remove base_dir on nominal __exit__().
         node_cleanup_on_bad_exit:   remove base_dir on __exit__() via exception.
+
+    NOTE: attributes must not be callable or begin with __.
     """
 
     cache_initdb = True
-    cached_initdb_dir = None
+    _cached_initdb_dir = None
 
     cache_pg_config = True
 
@@ -31,12 +40,129 @@ class TestgresConfig:
     node_cleanup_on_good_exit = True
     node_cleanup_on_bad_exit = False
 
+    @property
+    def cached_initdb_dir(self):
+        return self._cached_initdb_dir
+
+    @cached_initdb_dir.setter
+    def cached_initdb_dir(self, value):
+        self._cached_initdb_dir = value
+
+        if value:
+            cached_initdb_dirs.add(value)
+
+    def __init__(self, **options):
+        self.update(options)
+
+    def __setitem__(self, key, value):
+        setattr(self, key, value)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __setattr__(self, name, value):
+        if name not in self.keys():
+            raise TypeError('Unknown option {}'.format(name))
+
+        super(GlobalConfig, self).__setattr__(name, value)
+
+    def keys(self):
+        keys = []
+
+        for key in dir(GlobalConfig):
+            if not key.startswith('__') and not callable(self[key]):
+                keys.append(key)
+
+        return keys
+
+    def items(self):
+        return ((key, self[key]) for key in self.keys())
+
+    def update(self, config):
+        for key, value in config.items():
+            self[key] = value
+
+        return self
+
+    def copy(self):
+        return copy.copy(self)
+
+
+# cached dirs to be removed
+cached_initdb_dirs = set()
+
+# default config object
+testgres_config = GlobalConfig()
+
+# NOTE: for compatibility
+TestgresConfig = testgres_config
+
+# stack of GlobalConfigs
+config_stack = []
+
+
+def rm_cached_initdb_dirs():
+    for d in cached_initdb_dirs:
+        shutil.rmtree(d, ignore_errors=True)
+
+
+def push_config(**options):
+    """
+    Permanently set custom GlobalConfig options
+    and put previous settings on top of stack.
+    """
+
+    # push current config to stack
+    config_stack.append(testgres_config.copy())
+
+    return testgres_config.update(options)
+
+
+def pop_config():
+    """
+    Set previous GlobalConfig options from stack.
+    """
+
+    if len(config_stack) == 0:
+        raise IndexError('Reached initial config')
+
+    # restore popped config
+    return testgres_config.update(config_stack.pop())
+
+
+@contextmanager
+def scoped_config(**options):
+    """
+    Temporarily set custom GlobalConfig options for this context.
+
+    Example:
+        >>> with scoped_config(cache_initdb=False):
+        ...     with get_new_node().init().start() as node:
+        ...         print(node.execute('select 1'))
+    """
+
+    try:
+        # set a new config with options
+        config = push_config(**options)
+
+        # return it
+        yield config
+    finally:
+        # restore previous config
+        pop_config()
+
 
 def configure_testgres(**options):
     """
-    Configure testgres.
-    Look at TestgresConfig to check what can be changed.
+    Adjust current global options.
+    Look at GlobalConfig to learn what can be set.
     """
 
-    for key, option in options.items():
-        setattr(TestgresConfig, key, option)
+    testgres_config.update(options)
+
+
+# NOTE: to be executed at exit()
+atexit.register(rm_cached_initdb_dirs)
+
+# NOTE: assign initial cached dir for initdb
+testgres_config.cached_initdb_dir = tempfile.mkdtemp()

testgres/consts.py

@@ -4,6 +4,7 @@
 DATA_DIR = "data"
 LOGS_DIR = "logs"
 
+# names for config files
 RECOVERY_CONF_FILE = "recovery.conf"
 PG_AUTO_CONF_FILE = "postgresql.auto.conf"
 PG_CONF_FILE = "postgresql.conf"

testgres/exceptions.py

@@ -9,11 +9,7 @@ class TestgresException(Exception):
 
 @six.python_2_unicode_compatible
 class ExecUtilException(TestgresException):
-    def __init__(self,
-                 message=None,
-                 command=None,
-                 exit_code=0,
-                 out=None):
+    def __init__(self, message=None, command=None, exit_code=0, out=None):
         super(ExecUtilException, self).__init__(message)
 
         self.message = message

testgres/node.py

@@ -13,7 +13,7 @@
 
 from .cache import cached_initdb
 
-from .config import TestgresConfig
+from .config import testgres_config
 
 from .connection import \
     NodeConnection, \
@@ -88,8 +88,8 @@ def __init__(self, name=None, port=None, base_dir=None):
         self.base_dir = base_dir
 
         # defaults for __exit__()
-        self.cleanup_on_good_exit = TestgresConfig.node_cleanup_on_good_exit
-        self.cleanup_on_bad_exit = TestgresConfig.node_cleanup_on_bad_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
         self.shutdown_max_attempts = 3
 
         # private
@@ -150,7 +150,7 @@ def _try_shutdown(self, max_attempts):
                 self.stop()
                 break    # OK
             except ExecUtilException:
-                pass     # one more time
+                pass    # one more time
             except Exception:
                 # TODO: probably kill stray instance
                 eprint('cannot stop node {}'.format(self.name))
@@ -204,7 +204,7 @@ def _prepare_dirs(self):
             os.makedirs(self.logs_dir)
 
     def _maybe_start_logger(self):
-        if TestgresConfig.use_python_logging:
+        if testgres_config.use_python_logging:
             # spawn new logger if it doesn't exist or is stopped
             if not self._logger or not self._logger.is_alive():
                 self._logger = TestgresLogger(self.name, self.pg_log_name)
@@ -223,7 +223,7 @@ def _collect_special_files(self):
             (os.path.join(self.data_dir, PG_AUTO_CONF_FILE), 0),
             (os.path.join(self.data_dir, RECOVERY_CONF_FILE), 0),
             (os.path.join(self.data_dir, HBA_CONF_FILE), 0),
-            (self.pg_log_name, TestgresConfig.error_log_lines)
+            (self.pg_log_name, testgres_config.error_log_lines)
         ]
 
         for f, num_lines in files:
@@ -248,7 +248,7 @@ def init(self,
              fsync=False,
              unix_sockets=True,
              allow_streaming=True,
-             initdb_params=[]):
+             initdb_params=None):
         """
         Perform initdb for this node.
 
@@ -266,8 +266,9 @@ def init(self,
         self._prepare_dirs()
 
         # initialize this PostgreSQL node
-        initdb_log = os.path.join(self.logs_dir, "initdb.log")
-        cached_initdb(self.data_dir, initdb_log, initdb_params)
+        cached_initdb(data_dir=self.data_dir,
+                      logfile=self.utils_log_name,
+                      params=initdb_params)
 
         # initialize default config files
         self.default_conf(fsync=fsync,
@@ -603,7 +604,7 @@ def cleanup(self, max_attempts=3):
         self._try_shutdown(max_attempts)
 
         # choose directory to be removed
-        if TestgresConfig.node_cleanup_full:
+        if testgres_config.node_cleanup_full:
             rm_dir = self.base_dir    # everything
         else:
             rm_dir = self.data_dir    # just data, save logs