Merge branch 'master' of github.com:postgrespro/testgres

Author: ildus

README.md

@@ -2,6 +2,8 @@
 [![codecov](https://codecov.io/gh/postgrespro/testgres/branch/master/graph/badge.svg)](https://codecov.io/gh/postgrespro/testgres)
 [![PyPI version](https://badge.fury.io/py/testgres.svg)](https://badge.fury.io/py/testgres)
 
+[Documentation](https://postgrespro.github.io/testgres/)
+
 # testgres
 
 PostgreSQL testing utility. Both Python 2.7 and 3.3+ are supported.

setup.py

@@ -21,7 +21,7 @@
     readme = f.read()
 
 setup(
-    version='1.7.0',
+    version='1.8.2',
     name='testgres',
     packages=['testgres'],
     description='Testing utility for PostgreSQL and its extensions',

testgres/__init__.py

@@ -10,8 +10,10 @@
 
 from .connection import \
     NodeConnection, \
+    DatabaseError, \
     InternalError, \
-    ProgrammingError
+    ProgrammingError, \
+    OperationalError
 
 from .exceptions import *
 from .enums import *

testgres/backup.py

@@ -143,8 +143,8 @@ def spawn_primary(self, name=None, destroy=True):
         base_dir = self._prepare_dir(destroy)
 
         # Build a new PostgresNode
-        from .node import PostgresNode
-        with clean_on_error(PostgresNode(name=name, base_dir=base_dir)) as node:
+        NodeClass = self.original_node.__class__
+        with clean_on_error(NodeClass(name=name, base_dir=base_dir)) as node:
 
             # New nodes should always remove dir tree
             node._should_rm_dirs = True

testgres/connection.py

@@ -17,9 +17,11 @@
 
 from .exceptions import QueryException
 
-# export these exceptions
+# export some exceptions
+DatabaseError = pglib.DatabaseError
 InternalError = pglib.InternalError
 ProgrammingError = pglib.ProgrammingError
+OperationalError = pglib.OperationalError
 
 
 class NodeConnection(object):

testgres/node.py

@@ -19,10 +19,7 @@
 
 from .config import testgres_config
 
-from .connection import \
-    NodeConnection, \
-    InternalError,  \
-    ProgrammingError
+from .connection import NodeConnection
 
 from .consts import \
     DATA_DIR, \
@@ -541,10 +538,10 @@ def append_conf(self, line='', filename=PG_CONF_FILE, **kwargs):
             This instance of :class:`.PostgresNode`.
 
         Examples:
-            append_conf(fsync=False)
-            append_conf('log_connections = yes')
-            append_conf(random_page_cost=1.5, fsync=True, ...)
-            append_conf('postgresql.conf', 'synchronous_commit = off')
+            >>> append_conf(fsync=False)
+            >>> append_conf('log_connections = yes')
+            >>> append_conf(random_page_cost=1.5, fsync=True, ...)
+            >>> append_conf('postgresql.conf', 'synchronous_commit = off')
         """
 
         lines = [line]
@@ -815,7 +812,8 @@ def psql(self,
              filename=None,
              dbname=None,
              username=None,
-             input=None):
+             input=None,
+             **variables):
         """
         Execute a query using psql.
 
@@ -825,9 +823,15 @@ def psql(self,
             dbname: database name to connect to.
             username: database user name.
             input: raw input to be passed.
+            **variables: vars to be set before execution.
 
         Returns:
             A tuple of (code, stdout, stderr).
+
+        Examples:
+            >>> psql('select 1')
+            >>> psql('postgres', 'select 2')
+            >>> psql(query='select 3', ON_ERROR_STOP=1)
         """
 
         # Set default arguments
@@ -846,6 +850,10 @@ def psql(self,
             dbname
         ]  # yapf: disable
 
+        # set variables before execution
+        for key, value in iteritems(variables):
+            psql_params.extend(["--set", '{}={}'.format(key, value)])
+
         # select query source
         if query:
             psql_params.extend(("-c", query))
@@ -877,10 +885,15 @@ def safe_psql(self, query=None, **kwargs):
             username: database user name.
             input: raw input to be passed.
 
+            **kwargs are passed to psql().
+
         Returns:
             psql's output as str.
         """
 
+        # force this setting
+        kwargs['ON_ERROR_STOP'] = 1
+
         ret, out, err = self.psql(query=query, **kwargs)
         if ret:
             raise QueryException((err or b'').decode('utf-8'), query)
@@ -980,8 +993,7 @@ def poll_query_until(self,
                          sleep_time=1,
                          expected=True,
                          commit=True,
-                         raise_programming_error=True,
-                         raise_internal_error=True):
+                         suppress=None):
         """
         Run a query once per second until it returns 'expected'.
         Query should return a single value (1 row, 1 column).
@@ -994,13 +1006,13 @@ def poll_query_until(self,
             sleep_time: how much should we sleep after a failure?
             expected: what should be returned to break the cycle?
             commit: should (possible) changes be committed?
-            raise_programming_error: enable ProgrammingError?
-            raise_internal_error: enable InternalError?
+            suppress: a collection of exceptions to be suppressed.
 
         Examples:
-            poll_query_until('select true')
-            poll_query_until('postgres', "select now() > '01.01.2018'")
-            poll_query_until('select false', expected=True, max_attempts=4)
+            >>> poll_query_until('select true')
+            >>> poll_query_until('postgres', "select now() > '01.01.2018'")
+            >>> poll_query_until('select false', expected=True, max_attempts=4)
+            >>> poll_query_until('select 1', suppress={testgres.OperationalError})
         """
 
         # sanity checks
@@ -1032,13 +1044,8 @@ def poll_query_until(self,
                 elif expected is None:
                     return    # done
 
-            except ProgrammingError as e:
-                if raise_programming_error:
-                    raise e
-
-            except InternalError as e:
-                if raise_internal_error:
-                    raise e
+            except tuple(suppress or []):
+                pass    # we're suppressing them
 
             time.sleep(sleep_time)
             attempts += 1
@@ -1229,13 +1236,14 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
             options: additional options for pgbench (list).
 
             **kwargs: named options for pgbench.
-                Examples:
-                    pgbench_run(initialize=True, scale=2)
-                    pgbench_run(time=10)
                 Run pgbench --help to learn more.
 
         Returns:
             Stdout produced by pgbench.
+
+        Examples:
+            >>> pgbench_run(initialize=True, scale=2)
+            >>> pgbench_run(time=10)
         """
 
         # Set default arguments

tests/test_simple.py

@@ -626,7 +626,7 @@ def test_poll_query_until(self):
                     query='dummy2',
                     max_attempts=3,
                     sleep_time=0.01,
-                    raise_programming_error=False)
+                    suppress={testgres.ProgrammingError})
 
             # check 1 arg, ok
             node.poll_query_until('select true')
@@ -862,7 +862,7 @@ def test_version_management(self):
         with get_new_node() as node:
             self.assertTrue(isinstance(version, six.string_types))
             self.assertTrue(isinstance(node.version, PgVer))
-            self.assertTrue(node.version == version)
+            self.assertEqual(node.version, str(version))
 
     def test_child_pids(self):
         master_processes = [