Merge branch 'master' into gh-pages

funbringer · funbringer · commit f94931385fd0 · 2018-07-19T16:46:13.000+03:00
diff --git a/README.md b/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.
diff --git a/setup.py b/setup.py
@@ -21,7 +21,7 @@
     readme = f.read()
 
 setup(
-    version='1.8.0',
+    version='1.8.2',
     name='testgres',
     packages=['testgres'],
     description='Testing utility for PostgreSQL and its extensions',
diff --git a/testgres/backup.py b/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
diff --git a/testgres/node.py b/testgres/node.py
@@ -812,7 +812,8 @@ def psql(self,
              filename=None,
              dbname=None,
              username=None,
-             input=None):
+             input=None,
+             **variables):
         """
         Execute a query using psql.
 
@@ -822,9 +823,18 @@ 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')
+            (0, b'1\n', b'')
+            >>> psql('postgres', 'select 2')
+            (0, b'2\n', b'')
+            >>> psql(query='select 3', ON_ERROR_STOP=1)
+            (0, b'3\n', b'')
         """
 
         # Set default arguments
@@ -843,6 +853,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))
@@ -874,10 +888,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)
