minor refactoring

Author: zilder

testgres/connection.py

@@ -27,7 +27,12 @@ class NodeConnection(object):
     Transaction wrapper returned by Node
     """
 
-    def __init__(self, node, dbname=None, username=None, password=None):
+    def __init__(self,
+                 node,
+                 dbname=None,
+                 username=None,
+                 password=None,
+                 autocommit=False):
 
         # Set default arguments
         dbname = dbname or default_dbname()
@@ -42,6 +47,7 @@ def __init__(self, node, dbname=None, username=None, password=None):
             host=node.host,
             port=node.port)
 
+        self._connection.autocommit = autocommit
         self._cursor = self.connection.cursor()
 
     @property

testgres/consts.py

@@ -29,3 +29,6 @@
 MAX_REPLICATION_SLOTS = 10
 MAX_WAL_SENDERS = 10
 WAL_KEEP_SEGMENTS = 20
+
+# logical replication settings
+LOGICAL_REPL_MAX_CATCHUP_ATTEMPTS = 60

testgres/node.py

@@ -953,13 +953,11 @@ def execute(self,
 
         with self.connect(dbname=dbname,
                           username=username,
-                          password=password) as node_con:  # yapf: disable
+                          password=password,
+                          autocommit=commit) as node_con:  # yapf: disable
 
             res = node_con.execute(query)
 
-            if commit:
-                node_con.commit()
-
             return res
 
     def backup(self, **kwargs):
@@ -1152,14 +1150,21 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
 
         return execute_utility(_params, self.utils_log_file)
 
-    def connect(self, dbname=None, username=None, password=None):
+    def connect(self,
+                dbname=None,
+                username=None,
+                password=None,
+                autocommit=False):
         """
         Connect to a database.
 
         Args:
             dbname: database name to connect to.
             username: database user name.
             password: user's password.
+            autocommit: commit each statement automatically. Also it should be
+                set to `True` for statements requiring to be run outside
+                a transaction? such as `VACUUM` or `CREATE DATABASE`.
 
         Returns:
             An instance of :class:`.NodeConnection`.
@@ -1168,4 +1173,5 @@ def connect(self, dbname=None, username=None, password=None):
         return NodeConnection(node=self,
                               dbname=dbname,
                               username=username,
-                              password=password)  # yapf: disable
+                              password=password,
+                              autocommit=autocommit)  # yapf: disable

testgres/pubsub.py

@@ -12,7 +12,7 @@
 After that :meth:`~.PostgresNode.publish()` and :meth:`~.PostgresNode.subscribe()`
 methods may be used to setup replication. Example:
 
->>> from .api import get_new_node
+>>> from testgres import get_new_node
 >>> with get_new_node() as nodeA, get_new_node() as nodeB:
 ...     nodeA.init(allow_logical=True).start()
 ...     nodeB.init().start()
@@ -44,6 +44,7 @@
 
 from six import raise_from
 
+from .consts import LOGICAL_REPL_MAX_CATCHUP_ATTEMPTS
 from .defaults import default_dbname, default_username
 from .exceptions import CatchUpException
 from .utils import options_string
@@ -56,11 +57,11 @@ def __init__(self, name, node, tables=None, dbname=None, username=None):
         constructing publication objects.
 
         Args:
-            name: publication name
-            node: publisher's node
-            tables: tables list or None for all tables
-            dbname: database name used to connect and perform subscription
-            username: username used to connect to the database
+            name: publication name.
+            node: publisher's node.
+            tables: tables list or None for all tables.
+            dbname: database name used to connect and perform subscription.
+            username: username used to connect to the database.
         """
         self.name = name
         self.node = node
@@ -70,7 +71,7 @@ def __init__(self, name, node, tables=None, dbname=None, username=None):
         # create publication in database
         t = "table " + ", ".join(tables) if tables else "all tables"
         query = "create publication {} for {}"
-        node.safe_psql(query.format(name, t), dbname=dbname, username=username)
+        node.execute(query.format(name, t), dbname=dbname, username=username)
 
     def drop(self, dbname=None, username=None):
         """
@@ -87,13 +88,13 @@ def add_tables(self, tables, dbname=None, username=None):
         created with empty tables list.
 
         Args:
-            tables: a list of tables to be added to the publication
+            tables: a list of tables to be added to the publication.
         """
         if not tables:
             raise ValueError("Tables list is empty")
 
         query = "alter publication {} add table {}"
-        self.node.safe_psql(
+        self.node.execute(
             query.format(self.name, ", ".join(tables)),
             dbname=dbname or self.dbname,
             username=username or self.username)
@@ -112,15 +113,15 @@ def __init__(self,
         constructing subscription objects.
 
         Args:
-            name: subscription name
-            node: subscriber's node
+            name: subscription name.
+            node: subscriber's node.
             publication: :class:`.Publication` object we are subscribing to
-                (see :meth:`.PostgresNode.publish()`)
-            dbname: database name used to connect and perform subscription
-            username: username used to connect to the database
+                (see :meth:`.PostgresNode.publish()`).
+            dbname: database name used to connect and perform subscription.
+            username: username used to connect to the database.
             params: subscription parameters (see documentation on `CREATE SUBSCRIPTION
                  <https://www.postgresql.org/docs/current/static/sql-createsubscription.html>`_
-                 for details)
+                 for details).
         """
         self.name = name
         self.node = node
@@ -142,28 +143,29 @@ def __init__(self,
         if params:
             query += " with ({})".format(options_string(**params))
 
-        node.safe_psql(query, dbname=dbname, username=username)
+        # Note: cannot run 'create subscription' query in transaction mode
+        node.execute(query, dbname=dbname, username=username)
 
     def disable(self, dbname=None, username=None):
         """
         Disables the running subscription.
         """
         query = "alter subscription {} disable"
-        self.node.safe_psql(query.format(self.name), dbname=None, username=None)
+        self.node.execute(query.format(self.name), dbname=None, username=None)
 
     def enable(self, dbname=None, username=None):
         """
         Enables the previously disabled subscription.
         """
         query = "alter subscription {} enable"
-        self.node.safe_psql(query.format(self.name), dbname=None, username=None)
+        self.node.execute(query.format(self.name), dbname=None, username=None)
 
     def refresh(self, copy_data=True, dbname=None, username=None):
         """
         Disables the running subscription.
         """
         query = "alter subscription {} refresh publication with (copy_data={})"
-        self.node.safe_psql(
+        self.node.execute(
             query.format(self.name, copy_data),
             dbname=dbname,
             username=username)
@@ -172,7 +174,7 @@ def drop(self, dbname=None, username=None):
         """
         Drops subscription
         """
-        self.node.safe_psql(
+        self.node.execute(
             "drop subscription {}".format(self.name),
             dbname=dbname,
             username=username)
@@ -182,19 +184,19 @@ def catchup(self, username=None):
         Wait until subscription catches up with publication.
 
         Args:
-            username: remote node's user name
+            username: remote node's user name.
         """
-        query = (
-            "select pg_current_wal_lsn() - replay_lsn = 0 "
-            "from pg_stat_replication where application_name = '{}'").format(
-                self.name)
+        query = """
+            select pg_current_wal_lsn() - replay_lsn = 0
+            from pg_catalog.pg_stat_replication where application_name = '{}'
+        """.format(self.name)
 
         try:
             # wait until this LSN reaches subscriber
             self.pub.node.poll_query_until(
                 query=query,
                 dbname=self.pub.dbname,
                 username=username or self.pub.username,
-                max_attempts=60)
+                max_attempts=LOGICAL_REPL_MAX_CATCHUP_ATTEMPTS)
         except Exception as e:
             raise_from(CatchUpException("Failed to catch up", query), e)

tests/test_simple.py

@@ -432,7 +432,7 @@ def test_logical_replication(self):
             node1.safe_psql('insert into test2 values (\'a\'), (\'b\')')
             sub.catchup()
             res = node2.execute('select * from test2')
-            self.assertListEqual(res, [('a',), ('b',)])
+            self.assertListEqual(res, [('a', ), ('b', )])
 
             # drop subscription
             sub.drop()
@@ -450,12 +450,12 @@ def test_logical_replication(self):
 
             # explicitely add table
             with self.assertRaises(ValueError):
-                pub.add_tables([])  # fail
+                pub.add_tables([])    # fail
             pub.add_tables(['test2'])
             node1.safe_psql('insert into test2 values (\'c\')')
             sub.catchup()
             res = node2.execute('select * from test2')
-            self.assertListEqual(res, [('a',), ('b',)])
+            self.assertListEqual(res, [('a', ), ('b', )])
 
     @unittest.skipUnless(pg_version_ge('10'), 'requires 10+')
     def test_logical_catchup(self):
@@ -477,7 +477,10 @@ def test_logical_catchup(self):
                 node1.execute('insert into test values ({0}, {0})'.format(i))
                 sub.catchup()
                 res = node2.execute('select * from test')
-                self.assertListEqual(res, [(i, i,)])
+                self.assertListEqual(res, [(
+                    i,
+                    i,
+                )])
                 node1.execute('delete from test')
 
     @unittest.skipIf(pg_version_ge('10'), 'requires <10')
@@ -544,7 +547,8 @@ def test_poll_query_until(self):
 
             # check 0 columns
             with self.assertRaises(QueryException):
-                node.poll_query_until(query='select from pg_class limit 1')
+                node.poll_query_until(
+                    query='select from pg_catalog.pg_class limit 1')
 
             # check None, fail
             with self.assertRaises(QueryException):
@@ -556,7 +560,8 @@ def test_poll_query_until(self):
 
             # check 0 rows equivalent to expected=None
             node.poll_query_until(
-                query='select * from pg_class where true = false', expected=None)
+                query='select * from pg_catalog.pg_class where true = false',
+                expected=None)
 
             # check arbitrary expected value, fail
             with self.assertRaises(TimeoutException):