Add support for streaming replication protocol

doc/src/extras.rst

@@ -141,6 +141,128 @@ Logging cursor
 
 .. autoclass:: MinTimeLoggingCursor
 
+Replication cursor
+^^^^^^^^^^^^^^^^^^
+
+.. autoclass:: ReplicationConnection
+
+   This connection factory class can be used to open a special type of
+   connection that is used for streaming replication.
+
+   Example::
+
+       from psycopg2.extras import ReplicationConnection, REPLICATION_PHYSICAL, REPLICATION_LOGICAL
+       conn = psycopg2.connect(dsn, connection_factory=ReplicationConnection)
+       cur = conn.cursor()
+
+.. seealso::
+
+   - PostgreSQL `Replication protocol`__
+
+   .. __: http://www.postgresql.org/docs/current/static/protocol-replication.html
+
+.. autoclass:: ReplicationCursor
+
+    .. method:: identify_system()
+
+       Get information about the cluster status in form of a dict with
+       ``systemid``, ``timeline``, ``xlogpos`` and ``dbname`` as keys.
+
+       Example::
+
+           >>> print cur.identify_system()
+           {'timeline': 1, 'systemid': '1234567890123456789', 'dbname': 'test', 'xlogpos': '0/1ABCDEF'}
+
+    .. method:: create_replication_slot(slot_type, slot_name, output_plugin=None)
+
+       Create streaming replication slot.
+
+       :param slot_type: type of replication: either `REPLICATION_PHYSICAL` or
+                         `REPLICATION_LOGICAL`
+       :param slot_name: name of the replication slot to be created
+       :param output_plugin: name of the logical decoding output plugin to use
+                             (logical replication only)
+
+       Example::
+
+           cur.create_replication_slot(REPLICATION_LOGICAL, "testslot", "test_decoding")
+
+    .. method:: drop_replication_slot(slot_name)
+
+       Drop streaming replication slot.
+
+       :param slot_name: name of the replication slot to drop
+
+       Example::
+
+           cur.drop_replication_slot("testslot")
+
+    .. method:: start_replication(file, slot_type, slot_name=None, start_lsn=None, timeline=0, keepalive_interval=10, options=None)
+
+       Start and consume replication stream.
+
+       :param file: a file-like object to write replication stream messages to
+       :param slot_type: type of replication: either `REPLICATION_PHYSICAL` or
+                         `REPLICATION_LOGICAL`
+       :param slot_name: name of the replication slot to use (required for
+                         logical replication)
+       :param start_lsn: the point in replication stream (WAL position) to start
+                         from, in the form ``XXX/XXX`` (forward-slash separated
+                         pair of hexadecimals)
+       :param timeline: WAL history timeline to start streaming from (optional,
+                        can only be used with physical replication)
+       :param keepalive_interval: interval (in seconds) to send keepalive
+                                  messages to the server, in case there was no
+                                  communication during that period of time
+       :param options: an dictionary of options to pass to logical replication
+                       slot
+
+       The ``keepalive_interval`` must be greater than zero.
+
+       This method never returns unless an error message is sent from the
+       server, or the server closes connection, or there is an exception in the
+       ``write()`` method of the ``file`` object.
+
+       One can even use ``sys.stdout`` as the destination (this is only good for
+       testing purposes, however)::
+
+           >>> cur.start_replication(sys.stdout, "testslot")
+           ...
+
+       This method acts much like the `~cursor.copy_to()` with an important
+       distinction that ``write()`` method return value is dirving the
+       server-side replication cursor.  In order to report to the server that
+       the all the messages up to the current one have been stored reliably, one
+       should return true value (i.e. something that satisfies ``if retval:``
+       conidtion) from the ``write`` callback::
+
+           class ReplicationStreamWriter(object):
+               def write(self, msg):
+                   if store_message_reliably(msg):
+                       return True
+
+           cur.start_replication(writer, "testslot")
+           ...
+
+       .. note::
+
+           One needs to be aware that failure to update the server-side cursor
+           on any one replication slot properly by constantly consuming and
+           reporting success to the server can eventually lead to "disk full"
+           condition on the server, because the server retains all the WAL
+           segments that might be needed to stream the changes via currently
+           open replication slots.
+
+           Drop any open replication slots that are no longer being used.  The
+           list of open slots can be obtained by running a query like ``SELECT *
+           FROM pg_replication_slots``.
+
+.. data:: REPLICATION_PHYSICAL
+
+.. data:: REPLICATION_LOGICAL
+
+.. index::
+    pair: Cursor; Replication
 
 
 .. index::

lib/extras.py

@@ -437,6 +437,144 @@ def callproc(self, procname, vars=None):
         return LoggingCursor.callproc(self, procname, vars)
 
 
+class ReplicationConnection(_connection):
+    """A connection that uses `ReplicationCursor` automatically."""
+
+    def __init__(self, *args, **kwargs):
+        """Initializes a replication connection, by adding appropriate replication parameter to the provided dsn arguments."""
+
+        if len(args):
+            dsn = args[0]
+
+            # FIXME: could really use parse_dsn here
+
+            if dsn.startswith('postgres://') or dsn.startswith('postgresql://'):
+                # poor man's url parsing
+                if dsn.rfind('?') > 0:
+                    if not dsn.endswith('?'):
+                        dsn += '&'
+                else:
+                    dsn += '?'
+            else:
+                dsn += ' '
+            dsn += 'replication=database'
+            args = [dsn] + list(args[1:])
+        else:
+            dbname = kwargs.get('dbname', None)
+            if dbname is None:
+                kwargs['dbname'] = 'replication'
+
+            if kwargs.get('replication', None) is None:
+                kwargs['replication'] = 'database' if dbname else 'true'
+
+        super(ReplicationConnection, self).__init__(*args, **kwargs)
+
+        # prevent auto-issued BEGIN statements
+        self.autocommit = True
+
+    def cursor(self, *args, **kwargs):
+        kwargs.setdefault('cursor_factory', ReplicationCursor)
+        return super(ReplicationConnection, self).cursor(*args, **kwargs)
+
+
+"""Streamging replication types."""
+REPLICATION_PHYSICAL = 0
+REPLICATION_LOGICAL  = 1
+
+class ReplicationCursor(_cursor):
+    """A cursor used for replication commands."""
+
+    def identify_system(self):
+        """Get information about the cluster status."""
+
+        self.execute("IDENTIFY_SYSTEM")
+        return dict(zip(['systemid', 'timeline', 'xlogpos', 'dbname'],
+                        self.fetchall()[0]))
+
+    def quote_ident(self, ident):
+        # FIXME: use PQescapeIdentifier or psycopg_escape_identifier_easy, somehow
+        return '"%s"' % ident.replace('"', '""')
+
+    def create_replication_slot(self, slot_type, slot_name, output_plugin=None):
+        """Create streaming replication slot."""
+
+        command = "CREATE_REPLICATION_SLOT %s " % self.quote_ident(slot_name)
+
+        if slot_type == REPLICATION_LOGICAL:
+            if output_plugin is None:
+                raise RuntimeError("output_plugin is required for logical replication slot")
+
+            command += "LOGICAL %s" % self.quote_ident(output_plugin)
+
+        elif slot_type == REPLICATION_PHYSICAL:
+            if output_plugin is not None:
+                raise RuntimeError("output_plugin is not applicable to physical replication")
+
+            command += "PHYSICAL"
+
+        else:
+            raise RuntimeError("unrecognized replication slot type")
+
+        return self.execute(command)
+
+    def drop_replication_slot(self, slot_name):
+        """Drop streaming replication slot."""
+
+        command = "DROP_REPLICATION_SLOT %s" % self.quote_ident(slot_name)
+        return self.execute(command)
+
+    def start_replication(self, o, slot_type, slot_name=None, start_lsn=None,
+                          timeline=0, keepalive_interval=10, options=None):
+        """Start and consume replication stream."""
+
+        if keepalive_interval <= 0:
+            raise RuntimeError("keepalive_interval must be > 0: %d" % keepalive_interval)
+
+        command = "START_REPLICATION "
+
+        if slot_type == REPLICATION_LOGICAL and slot_name is None:
+                raise RuntimeError("slot_name is required for logical replication slot")
+
+        if slot_name:
+            command += "SLOT %s " % self.quote_ident(slot_name)
+
+        if slot_type == REPLICATION_LOGICAL:
+            command += "LOGICAL "
+        elif slot_type == REPLICATION_PHYSICAL:
+            command += "PHYSICAL "
+        else:
+            raise RuntimeError("unrecognized replication slot type")
+
+        if start_lsn is None:
+            start_lsn = '0/0'
+
+        # reparse lsn to catch possible garbage
+        lsn = start_lsn.split('/')
+        command += "%X/%X" % (int(lsn[0], 16), int(lsn[1], 16))
+
+        if timeline != 0:
+            if slot_type == REPLICATION_LOGICAL:
+                raise RuntimeError("cannot specify timeline for logical replication")
+
+            if timeline < 0:
+                raise RuntimeError("timeline must be >= 0: %d" % timeline)
+
+            command += " TIMELINE %d" % timeline
+
+        if options:
+            if slot_type == REPLICATION_PHYSICAL:
+                raise RuntimeError("cannot specify plugin options for physical replication")
+
+            command += " ("
+            for k,v in options.iteritems():
+                if not command.endswith('('):
+                    command += ", "
+                command += "%s %s" % (self.quote_ident(k), _A(str(v)).getquoted())
+            command += ")"
+
+        return self.start_replication_expert(o, command, keepalive_interval)
+
+
 # a dbtype and adapter for Python UUID type
 
 class UUID_adapter(object):

psycopg/cursor.h

@@ -72,6 +72,8 @@ struct cursorObject {
 #define DEFAULT_COPYSIZE 16384
 #define DEFAULT_COPYBUFF  8192
 
+    int   keepalive_interval;   /* interval for keepalive messages in replication mode */
+
     PyObject *tuple_factory;    /* factory for result tuples */
     PyObject *tzinfo_factory;   /* factory for tzinfo objects */
 
@@ -88,6 +90,10 @@ struct cursorObject {
 
 };
 
+/* streaming replication modes */
+#define CURSOR_REPLICATION_PHYSICAL 0
+#define CURSOR_REPLICATION_LOGICAL  1
+
 
 /* C-callable functions in cursor_int.c and cursor_type.c */
 BORROWED HIDDEN PyObject *curs_get_cast(cursorObject *self, PyObject *oid);

psycopg/cursor_type.c

@@ -1579,6 +1579,43 @@ psyco_curs_copy_expert(cursorObject *self, PyObject *args, PyObject *kwargs)
     return res;
 }
 
+#define psyco_curs_start_replication_expert_doc \
+"start_replication_expert(file, command, keepalive_interval) -- Start and consume replication stream with direct command."
+
+static PyObject *
+psyco_curs_start_replication_expert(cursorObject *self, PyObject *args)
+{
+    PyObject *file, *res = NULL;
+    char *command;
+    int keepalive_interval;
+
+    if (!PyArg_ParseTuple(args, "O&si",
+                          _psyco_curs_has_write_check, &file,
+                          &command, &keepalive_interval)) {
+        return NULL;
+    }
+
+    EXC_IF_CURS_CLOSED(self);
+    EXC_IF_CURS_ASYNC(self, start_replication_expert);
+    EXC_IF_GREEN(start_replication_expert);
+    EXC_IF_TPC_PREPARED(self->conn, start_replication_expert);
+
+    Dprintf("psyco_curs_start_replication_expert: command = %s", command);
+
+    self->copysize = 0;
+    Py_INCREF(file);
+    self->copyfile = file;
+    self->keepalive_interval = keepalive_interval;
+
+    if (pq_execute(self, command, 0, 1 /* no_result */, 1 /* no_begin */) >= 0) {
+        res = Py_None;
+        Py_INCREF(Py_None);
+    }
+    Py_CLEAR(self->copyfile);
+
+    return res;
+}
+
 /* extension: closed - return true if cursor is closed */
 
 #define psyco_curs_closed_doc \
@@ -1753,6 +1790,8 @@ static struct PyMethodDef cursorObject_methods[] = {
      METH_VARARGS|METH_KEYWORDS, psyco_curs_copy_to_doc},
     {"copy_expert", (PyCFunction)psyco_curs_copy_expert,
      METH_VARARGS|METH_KEYWORDS, psyco_curs_copy_expert_doc},
+    {"start_replication_expert", (PyCFunction)psyco_curs_start_replication_expert,
+     METH_VARARGS, psyco_curs_start_replication_expert_doc},
     {NULL}
 };