Merge branch 'master' into logical

zilder · zilder · commit d60cdcb753c4 · 2018-05-31T17:35:51.000+03:00
diff --git a/hooks/pre-commit b/hooks/pre-commit
@@ -1,13 +1,23 @@
 #!/bin/bash
 
+set -e
+
 # capture the changed files that have been staged
 changed_files=$(git diff --staged --name-only)
 
 for file in ${changed_files}
 do
 	if [[ "${file##*.}" == "py" ]]; then
-		echo "Yapfing ${file}"
-		yapf ${file} -i
-		git add ${file}
+		if command -v yapf > /dev/null; then
+			echo "Run yapf on ${file}"
+			yapf ${file} -i
+			git add ${file}
+		fi
+
+		if command -v flake8 > /dev/null; then
+			echo "Run flake8 on ${file}"
+			flake8 ${file}
+		fi
 	fi
 done
+
diff --git a/setup.py b/setup.py
@@ -16,12 +16,18 @@
 if sys.version_info < (3, 3):
     install_requires.append("ipaddress")
 
+# Get contents of README file
+with open('README.md', 'r') as f:
+    readme = f.read()
+
 setup(
-    version='1.6.0',
+    version='1.7.0',
     name='testgres',
     packages=['testgres'],
     description='Testing utility for PostgreSQL and its extensions',
     url='https://github.com/postgrespro/testgres',
+    long_description=readme,
+    long_description_content_type='text/markdown',
     license='PostgreSQL',
     author='Ildar Musin',
     author_email='zildermann@gmail.com',
diff --git a/testgres/enums.py b/testgres/enums.py
@@ -85,3 +85,14 @@ def from_process(process):
 
         # default
         return ProcessType.Unknown
+
+
+class DumpFormat(Enum):
+    """
+    Available dump formats
+    """
+
+    Plain = 'plain'
+    Custom = 'custom'
+    Directory = 'directory'
+    Tar = 'tar'
diff --git a/testgres/node.py b/testgres/node.py
@@ -10,7 +10,10 @@
 from six import raise_from, iteritems
 from tempfile import mkstemp, mkdtemp
 
-from .enums import NodeStatus, ProcessType
+from .enums import \
+    NodeStatus, \
+    ProcessType, \
+    DumpFormat
 
 from .cache import cached_initdb
 
@@ -55,7 +58,8 @@
     StartNodeException, \
     TimeoutException,   \
     InitNodeException,  \
-    TestgresException
+    TestgresException,  \
+    BackupException
 
 from .logger import TestgresLogger
 
@@ -577,12 +581,13 @@ def get_control_data(self):
 
         return out_dict
 
-    def start(self, params=[]):
+    def start(self, params=[], wait=True):
         """
         Start this node using pg_ctl.
 
         Args:
             params: additional arguments for pg_ctl.
+            wait: wait until operation completes.
 
         Returns:
             This instance of :class:`.PostgresNode`.
@@ -592,7 +597,7 @@ def start(self, params=[]):
             get_bin_path("pg_ctl"),
             "-D", self.data_dir,
             "-l", self.pg_log_file,
-            "-w",  # wait
+            "-w" if wait else '-W',  # --wait or --no-wait
             "start"
         ] + params  # yapf: disable
 
@@ -607,12 +612,13 @@ def start(self, params=[]):
 
         return self
 
-    def stop(self, params=[]):
+    def stop(self, params=[], wait=True):
         """
         Stop this node using pg_ctl.
 
         Args:
             params: additional arguments for pg_ctl.
+            wait: wait until operation completes.
 
         Returns:
             This instance of :class:`.PostgresNode`.
@@ -621,7 +627,7 @@ def stop(self, params=[]):
         _params = [
             get_bin_path("pg_ctl"),
             "-D", self.data_dir,
-            "-w",  # wait
+            "-w" if wait else '-W',  # --wait or --no-wait
             "stop"
         ] + params  # yapf: disable
 
@@ -681,6 +687,8 @@ def reload(self, params=[]):
 
         execute_utility(_params, self.utils_log_file)
 
+        return self
+
     def pg_ctl(self, params):
         """
         Invoke pg_ctl with params.
@@ -812,7 +820,11 @@ def safe_psql(self, query=None, **kwargs):
 
         return out
 
-    def dump(self, filename=None, dbname=None, username=None):
+    def dump(self,
+             filename=None,
+             dbname=None,
+             username=None,
+             format=DumpFormat.Plain):
         """
         Dump database into a file using pg_dump.
         NOTE: the file is not removed automatically.
@@ -821,14 +833,27 @@ def dump(self, filename=None, dbname=None, username=None):
             filename: database dump taken by pg_dump.
             dbname: database name to connect to.
             username: database user name.
+            format: format argument plain/custom/directory/tar.
 
         Returns:
             Path to a file containing dump.
         """
 
+        # Check arguments
+        if not isinstance(format, DumpFormat):
+            try:
+                format = DumpFormat(format)
+            except ValueError:
+                msg = 'Invalid format "{}"'.format(format)
+                raise BackupException(msg)
+
+        # Generate tmpfile or tmpdir
         def tmpfile():
-            fd, fname = mkstemp(prefix=TMP_DUMP)
-            os.close(fd)
+            if format == DumpFormat.Directory:
+                fname = mkdtemp(prefix=TMP_DUMP)
+            else:
+                fd, fname = mkstemp(prefix=TMP_DUMP)
+                os.close(fd)
             return fname
 
         # Set default arguments
@@ -842,7 +867,8 @@ def tmpfile():
             "-h", self.host,
             "-f", filename,
             "-U", username,
-            "-d", dbname
+            "-d", dbname,
+            "-F", format.value
         ]  # yapf: disable
 
         execute_utility(_params, self.utils_log_file)
@@ -854,12 +880,29 @@ def restore(self, filename, dbname=None, username=None):
         Restore database from pg_dump's file.
 
         Args:
-            filename: database dump taken by pg_dump.
+            filename: database dump taken by pg_dump in custom/directory/tar formats.
             dbname: database name to connect to.
             username: database user name.
         """
 
-        self.psql(filename=filename, dbname=dbname, username=username)
+        # Set default arguments
+        dbname = dbname or default_dbname()
+        username = username or default_username()
+
+        _params = [
+            get_bin_path("pg_restore"),
+            "-p", str(self.port),
+            "-h", self.host,
+            "-U", username,
+            "-d", dbname,
+            filename
+        ]  # yapf: disable
+
+        # try pg_restore if dump is binary formate, and psql if not
+        try:
+            execute_utility(_params, self.utils_log_name)
+        except ExecUtilException:
+            self.psql(filename=filename, dbname=dbname, username=username)
 
     @method_decorator(positional_args_hack(['dbname', 'query']))
     def poll_query_until(self,
diff --git a/testgres/pubsub.py b/testgres/pubsub.py
@@ -102,9 +102,9 @@ def add_tables(self, tables, dbname=None, username=None):
 
 class Subscription(object):
     def __init__(self,
-                 name,
                  node,
                  publication,
+                 name=None,
                  dbname=None,
                  username=None,
                  **params):
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -46,8 +46,8 @@
 
 def util_exists(util):
     def good_properties(f):
-        return (os.path.exists(f) and
-                os.path.isfile(f) and
+        return (os.path.exists(f) and  # noqa: W504
+                os.path.isfile(f) and  # noqa: W504
                 os.access(f, os.X_OK))  # yapf: disable
 
     # try to resolve it
@@ -67,12 +67,14 @@ def removing(f):
     finally:
         if os.path.isfile(f):
             os.remove(f)
+        elif os.path.isdir(f):
+            rmtree(f, ignore_errors=True)
 
 
 class TestgresTests(unittest.TestCase):
     def test_node_repr(self):
         with get_new_node() as node:
-            pattern = 'PostgresNode\(name=\'.+\', port=.+, base_dir=\'.+\'\)'
+            pattern = r"PostgresNode\(name='.+', port=.+, base_dir='.+'\)"
             self.assertIsNotNone(re.match(pattern, str(node)))
 
     def test_custom_init(self):
@@ -263,7 +265,7 @@ def test_psql(self):
             # check feeding input
             node.safe_psql('create table horns (w int)')
             node.safe_psql(
-                'copy horns from stdin (format csv)', input=b"1\n2\n3\n\.\n")
+                'copy horns from stdin (format csv)', input=b"1\n2\n3\n\\.\n")
             _sum = node.safe_psql('select sum(w) from horns')
             self.assertEqual(_sum, b'6\n')
 
@@ -373,6 +375,18 @@ def test_backup_wrong_xlog_method(self):
                     BackupException, msg='Invalid xlog_method "wrong"'):
                 node.backup(xlog_method='wrong')
 
+    def test_pg_ctl_wait_option(self):
+        with get_new_node() as node:
+            node.init().start(wait=False)
+            while True:
+                try:
+                    node.stop(wait=False)
+                    break
+                except ExecUtilException:
+                    # it's ok to get this exception here since node
+                    # could be not started yet
+                    pass
+
     def test_replicate(self):
         with get_new_node() as node:
             node.init(allow_streaming=True).start()
@@ -515,16 +529,17 @@ def test_dump(self):
         with get_new_node().init().start() as node1:
 
             node1.execute(query_create)
-
-            # take a new dump
-            with removing(node1.dump()) as dump:
-                with get_new_node().init().start() as node2:
-                    # restore dump
-                    self.assertTrue(os.path.isfile(dump))
-                    node2.restore(filename=dump)
-
-                    res = node2.execute(query_select)
-                    self.assertListEqual(res, [(1, ), (2, )])
+            for format in ['plain', 'custom', 'directory', 'tar']:
+                with removing(node1.dump(format=format)) as dump:
+                    with get_new_node().init().start() as node3:
+                        if format == 'directory':
+                            self.assertTrue(os.path.isdir(dump))
+                        else:
+                            self.assertTrue(os.path.isfile(dump))
+                        # restore dump
+                        node3.restore(filename=dump)
+                        res = node3.execute(query_select)
+                        self.assertListEqual(res, [(1, ), (2, )])
 
     def test_users(self):
         with get_new_node().init().start() as node:
