Logical replication support

testgres/node.py

@@ -50,10 +50,13 @@
     QueryException,     \
     StartNodeException, \
     TimeoutException,   \
+    InitNodeException,  \
     TestgresException
 
 from .logger import TestgresLogger
 
+from .pubsub import Publication, Subscription
+
 from .utils import \
     eprint, \
     get_bin_path, \
@@ -62,6 +65,7 @@
     reserve_port, \
     release_port, \
     execute_utility, \
+    options_string, \
     clean_on_error
 
 from .backup import NodeBackup
@@ -293,25 +297,25 @@ def _create_recovery_conf(self, username, slot=None):
         assert master is not None
 
         # yapf: disable
-        conninfo = (
-            u"application_name={} "
-            u"port={} "
-            u"user={} "
-        ).format(self.name, master.port, username)
+        conninfo = {
+            "application_name": self.name,
+            "port": master.port,
+            "user": username
+        }
 
         # host is tricky
         try:
             import ipaddress
             ipaddress.ip_address(master.host)
-            conninfo += u"hostaddr={}".format(master.host)
+            conninfo["hostaddr"] = master.host
         except ValueError:
-            conninfo += u"host={}".format(master.host)
+            conninfo["host"] = master.host
 
         # yapf: disable
         line = (
             "primary_conninfo='{}'\n"
             "standby_mode=on\n"
-        ).format(conninfo)
+        ).format(options_string(**conninfo))
 
         if slot:
             # Connect to master for some additional actions
@@ -405,6 +409,7 @@ def default_conf(self,
                      fsync=False,
                      unix_sockets=True,
                      allow_streaming=True,
+                     allow_logical=False,
                      log_statement='all'):
         """
         Apply default settings to this node.
@@ -413,6 +418,7 @@ def default_conf(self,
             fsync: should this node use fsync to keep data safe?
             unix_sockets: should we enable UNIX sockets?
             allow_streaming: should this node add a hba entry for replication?
+            allow_logical: can this node be used as a logical replication publisher?
             log_statement: one of ('all', 'off', 'mod', 'ddl').
 
         Returns:
@@ -494,6 +500,12 @@ def get_auth_method(t):
                                                       wal_keep_segments,
                                                       wal_level))
 
+            if allow_logical:
+                if not pg_version_ge('10'):
+                    raise InitNodeException("Logical replication is only "
+                                            "available for Postgres 10 and newer")
+                conf.write(u"wal_level = logical\n")
+
             # disable UNIX sockets if asked to
             if not unix_sockets:
                 conf.write(u"unix_socket_directories = ''\n")
@@ -903,13 +915,14 @@ def poll_query_until(self,
                 if res is None:
                     raise QueryException('Query returned None', query)
 
-                if len(res) == 0:
-                    raise QueryException('Query returned 0 rows', query)
-
-                if len(res[0]) == 0:
-                    raise QueryException('Query returned 0 columns', query)
-
-                if res[0][0] == expected:
+                # result set is not empty
+                if len(res):
+                    if len(res[0]) == 0:
+                        raise QueryException('Query returned 0 columns', query)
+                    if res[0][0] == expected:
+                        return    # done
+                # empty result set is considered as None
+                elif expected is None:
                     return    # done
 
             except ProgrammingError as e:
@@ -1018,6 +1031,40 @@ def catchup(self, dbname=None, username=None):
         except Exception as e:
             raise_from(CatchUpException("Failed to catch up", poll_lsn), e)
 
+    def publish(self,
+                name,
+                tables=None,
+                dbname=None,
+                username=None):
+        """
+        Create publication for logical replication
+
+        Args:
+            pubname: publication name
+            tables: tables names list
+            dbname: database name where objects or interest are located
+            username: replication username
+        """
+        return Publication(name=name, node=self, tables=tables, dbname=dbname,
+                           username=username)
+
+    def subscribe(self,
+                  publication,
+                  name,
+                  dbname=None,
+                  username=None,
+                  **kwargs):
+        """
+        Create subscription for logical replication
+
+        Args:
+            subname: subscription name
+            publication: publication object obtained from publish()
+
+        """
+        return Subscription(name=name, node=self, publication=publication,
+                            dbname=dbname, username=username, **kwargs)
+
     def pgbench(self,
                 dbname=None,
                 username=None,

testgres/pubsub.py

@@ -0,0 +1,153 @@
+# coding: utf-8
+
+from six import raise_from
+
+from .defaults import default_dbname, default_username
+from .exceptions import CatchUpException
+from .utils import options_string
+
+
+class Publication(object):
+    def __init__(self, name, node, tables=None, dbname=None, username=None):
+        """
+        Constructor
+
+        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
+        """
+        self.name = name
+        self.node = node
+        self.dbname = dbname or default_dbname()
+        self.username = username or default_username()
+
+        # 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)
+
+    def drop(self, dbname=None, username=None):
+        """
+        Drop publication
+        """
+        self.node.safe_psql(
+            "drop publication {}".format(self.name),
+            dbname=dbname,
+            username=username)
+
+    def add_tables(self, tables, dbname=None, username=None):
+        """
+        Add tables
+
+        Args:
+            tables: a list of tables to add to the publication
+        """
+        if not tables:
+            raise ValueError("Tables list is empty")
+
+        query = "alter publication {} add table {}"
+        self.node.safe_psql(
+            query.format(self.name, ", ".join(tables)),
+            dbname=dbname or self.dbname,
+            username=username or self.username)
+
+
+class Subscription(object):
+    def __init__(self,
+                 name,
+                 node,
+                 publication,
+                 dbname=None,
+                 username=None,
+                 **kwargs):
+        """
+        Constructor
+
+        Args:
+            name: subscription name
+            node: subscriber's node
+            publication: Publication object we are subscribing to
+            dbname: database name used to connect and perform subscription
+            username: username used to connect to the database
+            **kwargs: subscription parameters (see CREATE SUBSCRIPTION
+                in PostgreSQL documentation for more information)
+        """
+        self.name = name
+        self.node = node
+        self.pub = publication
+
+        # connection info
+        conninfo = {
+            "dbname": self.pub.dbname,
+            "user": self.pub.username,
+            "host": self.pub.node.host,
+            "port": self.pub.node.port
+        }
+
+        query = (
+            "create subscription {} connection '{}' publication {}").format(
+                name, options_string(**conninfo), self.pub.name)
+
+        # additional parameters
+        if kwargs:
+            query += " with ({})".format(options_string(**kwargs))
+
+        node.safe_psql(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)
+
+    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)
+
+    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(
+            query.format(self.name, copy_data),
+            dbname=dbname,
+            username=username)
+
+    def drop(self, dbname=None, username=None):
+        """
+        Drops subscription
+        """
+        self.node.safe_psql(
+            "drop subscription {}".format(self.name),
+            dbname=dbname,
+            username=username)
+
+    def catchup(self, username=None):
+        """
+        Wait until subscription catches up with publication.
+
+        Args:
+            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)
+
+        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)
+        except Exception as e:
+            raise_from(CatchUpException("Failed to catch up", query), e)

testgres/utils.py

@@ -11,6 +11,7 @@
 
 from contextlib import contextmanager
 from distutils.version import LooseVersion
+from six import iteritems
 
 from .config import testgres_config
 from .exceptions import ExecUtilException
@@ -225,6 +226,10 @@ def eprint(*args, **kwargs):
     print(*args, file=sys.stderr, **kwargs)
 
 
+def options_string(separator=u" ", **kwargs):
+    return separator.join(u"{}={}".format(k, v) for k, v in iteritems(kwargs))
+
+
 @contextmanager
 def clean_on_error(node):
     """