huge refactoring, split code into several submodules

Author: funbringer

MANIFEST

@@ -2,4 +2,12 @@
 setup.cfg
 setup.py
 testgres/__init__.py
-testgres/testgres.py
+testgres/api.py
+testgres/backup.py
+testgres/config.py
+testgres/connection.py
+testgres/consts.py
+testgres/exceptions.py
+testgres/logger.py
+testgres/node.py
+testgres/utils.py

testgres/__init__.py

@@ -1 +1,20 @@
-from .testgres import *
+from .api import get_new_node
+from .backup import NodeBackup
+from .config import TestgresConfig, configure_testgres
+
+from .connection import \
+    IsolationLevel, \
+    NodeConnection, \
+    InternalError, \
+    ProgrammingError
+
+from .exceptions import *
+from .node import NodeStatus, PostgresNode
+
+from .utils import \
+    reserve_port, \
+    release_port, \
+    bound_ports, \
+    get_bin_path, \
+    get_pg_config, \
+    get_pg_version

testgres/api.py

@@ -0,0 +1,45 @@
+# coding: utf-8
+"""
+Testing framework for PostgreSQL and its extensions
+
+This module was created under influence of Postgres TAP test feature
+(PostgresNode.pm module). It can manage Postgres clusters: initialize,
+edit configuration files, start/stop cluster, execute queries. The
+typical flow may look like:
+
+    with get_new_node('test') as node:
+        node.init()
+        node.start()
+        result = node.psql('postgres', 'SELECT 1')
+        print(result)
+        node.stop()
+
+    Or:
+
+    with get_new_node('node1') as node1:
+        node1.init().start()
+        with node1.backup() as backup:
+            with backup.spawn_primary('node2') as node2:
+                res = node2.start().execute('postgres', 'select 2')
+                print(res)
+
+Copyright (c) 2016, Postgres Professional
+"""
+
+from .node import PostgresNode
+
+
+def get_new_node(name, base_dir=None, use_logging=False):
+    """
+    Create a new node (select port automatically).
+
+    Args:
+        name: node's name.
+        base_dir: path to node's data directory.
+        use_logging: enable python logging.
+
+    Returns:
+        An instance of PostgresNode.
+    """
+
+    return PostgresNode(name=name, base_dir=base_dir, use_logging=use_logging)

testgres/backup.py

@@ -0,0 +1,152 @@
+# coding: utf-8
+
+import os
+import shutil
+import tempfile
+
+from .consts import \
+    DATA_DIR as _DATA_DIR, \
+    BACKUP_LOG_FILE as _BACKUP_LOG_FILE, \
+    DEFAULT_XLOG_METHOD as _DEFAULT_XLOG_METHOD
+
+from .exceptions import \
+    BackupException
+
+from .utils import \
+    default_username as _default_username, \
+    execute_utility as _execute_utility, \
+    explain_exception as _explain_exception
+
+
+class NodeBackup(object):
+    """
+    Smart object responsible for backups
+    """
+
+    @property
+    def log_file(self):
+        return os.path.join(self.base_dir, _BACKUP_LOG_FILE)
+
+    def __init__(self,
+                 node,
+                 base_dir=None,
+                 username=None,
+                 xlog_method=_DEFAULT_XLOG_METHOD):
+
+        if not node.status():
+            raise BackupException('Node must be running')
+
+        # Set default arguments
+        username = username or _default_username()
+        base_dir = base_dir or tempfile.mkdtemp()
+
+        # public
+        self.original_node = node
+        self.base_dir = base_dir
+
+        # private
+        self._available = True
+
+        data_dir = os.path.join(self.base_dir, _DATA_DIR)
+        _params = [
+            "-D{}".format(data_dir), "-p{}".format(node.port),
+            "-U{}".format(username), "-X{}".format(xlog_method)
+        ]
+        _execute_utility("pg_basebackup", _params, self.log_file)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.cleanup()
+
+    def _prepare_dir(self, destroy):
+        """
+        Provide a data directory for a copy of node.
+
+        Args:
+            destroy: should we convert this backup into a node?
+
+        Returns:
+            Path to data directory.
+        """
+
+        if not self._available:
+            raise BackupException('Backup is exhausted')
+
+        # Do we want to use this backup several times?
+        available = not destroy
+
+        if available:
+            dest_base_dir = tempfile.mkdtemp()
+
+            data1 = os.path.join(self.base_dir, _DATA_DIR)
+            data2 = os.path.join(dest_base_dir, _DATA_DIR)
+
+            try:
+                # Copy backup to new data dir
+                shutil.copytree(data1, data2)
+            except Exception as e:
+                raise BackupException(_explain_exception(e))
+        else:
+            dest_base_dir = self.base_dir
+
+        # Is this backup exhausted?
+        self._available = available
+
+        # Return path to new node
+        return dest_base_dir
+
+    def spawn_primary(self, name, destroy=True, use_logging=False):
+        """
+        Create a primary node from a backup.
+
+        Args:
+            name: name for a new node.
+            destroy: should we convert this backup into a node?
+            use_logging: enable python logging.
+
+        Returns:
+            New instance of PostgresNode.
+        """
+
+        base_dir = self._prepare_dir(destroy)
+
+        # Build a new PostgresNode
+        from .node import PostgresNode
+        node = PostgresNode(
+            name=name,
+            base_dir=base_dir,
+            master=self.original_node,
+            use_logging=use_logging)
+
+        # New nodes should always remove dir tree
+        node._should_rm_dirs = True
+
+        node.append_conf("postgresql.conf", "\n")
+        node.append_conf("postgresql.conf", "port = {}".format(node.port))
+
+        return node
+
+    def spawn_replica(self, name, destroy=True, use_logging=False):
+        """
+        Create a replica of the original node from a backup.
+
+        Args:
+            name: name for a new node.
+            destroy: should we convert this backup into a node?
+            use_logging: enable python logging.
+
+        Returns:
+            New instance of PostgresNode.
+        """
+
+        node = self.spawn_primary(name, destroy, use_logging=use_logging)
+        node._create_recovery_conf(self.original_node)
+
+        return node
+
+    def cleanup(self):
+        if self._available:
+            shutil.rmtree(self.base_dir, ignore_errors=True)
+            self._available = False

testgres/config.py

@@ -0,0 +1,29 @@
+# coding: utf-8
+
+
+class TestgresConfig:
+    """
+    Global config (override default settings).
+    """
+
+    # shall we cache pg_config results?
+    cache_pg_config = True
+
+    # shall we use cached initdb instance?
+    cache_initdb = True
+
+    # shall we create a temp dir for cached initdb?
+    cached_initdb_dir = None
+
+    # shall we remove EVERYTHING (including logs)?
+    node_cleanup_full = True
+
+
+def configure_testgres(**options):
+    """
+    Configure testgres.
+    Look at TestgresConfig to check what can be changed.
+    """
+
+    for key, option in options.items():
+        setattr(TestgresConfig, key, option)

testgres/connection.py

@@ -0,0 +1,125 @@
+# coding: utf-8
+
+# we support both pg8000 and psycopg2
+try:
+    import psycopg2 as pglib
+except ImportError:
+    try:
+        import pg8000 as pglib
+    except ImportError:
+        raise ImportError("You must have psycopg2 or pg8000 modules installed")
+
+from enum import Enum
+
+from .exceptions import QueryException
+from .utils import default_username as _default_username
+
+# export these exceptions
+InternalError = pglib.InternalError
+ProgrammingError = pglib.ProgrammingError
+
+
+class IsolationLevel(Enum):
+    """
+    Transaction isolation level for NodeConnection
+    """
+
+    ReadUncommitted, ReadCommitted, RepeatableRead, Serializable = range(4)
+
+
+class NodeConnection(object):
+    """
+    Transaction wrapper returned by Node
+    """
+
+    def __init__(self,
+                 parent_node,
+                 dbname,
+                 host="127.0.0.1",
+                 username=None,
+                 password=None):
+
+        # Use default user if not specified
+        username = username or _default_username()
+
+        self.parent_node = parent_node
+
+        self.connection = pglib.connect(
+            database=dbname,
+            user=username,
+            port=parent_node.port,
+            host=host,
+            password=password)
+
+        self.cursor = self.connection.cursor()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
+
+    def begin(self, isolation_level=IsolationLevel.ReadCommitted):
+        # yapf: disable
+        levels = [
+            'read uncommitted',
+            'read committed',
+            'repeatable read',
+            'serializable'
+        ]
+
+        # Check if level is an IsolationLevel
+        if (isinstance(isolation_level, IsolationLevel)):
+
+            # Get index of isolation level
+            level_idx = isolation_level.value
+            assert(level_idx in range(4))
+
+            # Replace isolation level with its name
+            isolation_level = levels[level_idx]
+
+        else:
+            # Get name of isolation level
+            level_str = str(isolation_level).lower()
+
+            # Validate level string
+            if level_str not in levels:
+                error = 'Invalid isolation level "{}"'
+                raise QueryException(error.format(level_str))
+
+            # Replace isolation level with its name
+            isolation_level = level_str
+
+        # Set isolation level
+        cmd = 'SET TRANSACTION ISOLATION LEVEL {}'
+        self.cursor.execute(cmd.format(isolation_level))
+
+        return self
+
+    def commit(self):
+        self.connection.commit()
+
+        return self
+
+    def rollback(self):
+        self.connection.rollback()
+
+        return self
+
+    def execute(self, query, *args):
+        self.cursor.execute(query, args)
+
+        try:
+            res = self.cursor.fetchall()
+
+            # pg8000 might return tuples
+            if isinstance(res, tuple):
+                res = [tuple(t) for t in res]
+
+            return res
+        except Exception:
+            return None
+
+    def close(self):
+        self.cursor.close()
+        self.connection.close()