Add support for streaming replication protocol

lib/extras.py

@@ -39,8 +39,12 @@
 from psycopg2 import extensions as _ext
 from psycopg2.extensions import cursor as _cursor
 from psycopg2.extensions import connection as _connection
-from psycopg2.extensions import adapt as _A
+from psycopg2.extensions import adapt as _A, quote_ident
 from psycopg2.extensions import b
+from psycopg2._psycopg import REPLICATION_PHYSICAL, REPLICATION_LOGICAL
+from psycopg2._psycopg import ReplicationConnection as _replicationConnection
+from psycopg2._psycopg import ReplicationCursor as _replicationCursor
+from psycopg2._psycopg import ReplicationMessage
 
 
 class DictCursorBase(_cursor):
@@ -437,6 +441,123 @@ def callproc(self, procname, vars=None):
         return LoggingCursor.callproc(self, procname, vars)
 
 
+class LogicalReplicationConnection(_replicationConnection):
+
+    def __init__(self, *args, **kwargs):
+        kwargs['replication_type'] = REPLICATION_LOGICAL
+        super(LogicalReplicationConnection, self).__init__(*args, **kwargs)
+
+
+class PhysicalReplicationConnection(_replicationConnection):
+
+    def __init__(self, *args, **kwargs):
+        kwargs['replication_type'] = REPLICATION_PHYSICAL
+        super(PhysicalReplicationConnection, self).__init__(*args, **kwargs)
+
+
+class StopReplication(Exception):
+    """
+    Exception used to break out of the endless loop in
+    `~ReplicationCursor.consume_stream()`.
+
+    Subclass of `~exceptions.Exception`.  Intentionally *not* inherited from
+    `~psycopg2.Error` as occurrence of this exception does not indicate an
+    error.
+    """
+    pass
+
+
+class ReplicationCursor(_replicationCursor):
+    """A cursor used for communication on replication connections."""
+
+    def create_replication_slot(self, slot_name, slot_type=None, output_plugin=None):
+        """Create streaming replication slot."""
+
+        command = "CREATE_REPLICATION_SLOT %s " % quote_ident(slot_name, self)
+
+        if slot_type is None:
+            slot_type = self.connection.replication_type
+
+        if slot_type == REPLICATION_LOGICAL:
+            if output_plugin is None:
+                raise psycopg2.ProgrammingError("output plugin name is required to create logical replication slot")
+
+            command += "LOGICAL %s" % quote_ident(output_plugin, self)
+
+        elif slot_type == REPLICATION_PHYSICAL:
+            if output_plugin is not None:
+                raise psycopg2.ProgrammingError("cannot specify output plugin name when creating physical replication slot")
+
+            command += "PHYSICAL"
+
+        else:
+            raise psycopg2.ProgrammingError("unrecognized replication type: %s" % repr(slot_type))
+
+        self.execute(command)
+
+    def drop_replication_slot(self, slot_name):
+        """Drop streaming replication slot."""
+
+        command = "DROP_REPLICATION_SLOT %s" % quote_ident(slot_name, self)
+        self.execute(command)
+
+    def start_replication(self, slot_name=None, slot_type=None, start_lsn=0,
+                          timeline=0, options=None, decode=False):
+        """Start replication stream."""
+
+        command = "START_REPLICATION "
+
+        if slot_type is None:
+            slot_type = self.connection.replication_type
+
+        if slot_type == REPLICATION_LOGICAL:
+            if slot_name:
+                command += "SLOT %s " % quote_ident(slot_name, self)
+            else:
+                raise psycopg2.ProgrammingError("slot name is required for logical replication")
+
+            command += "LOGICAL "
+
+        elif slot_type == REPLICATION_PHYSICAL:
+            if slot_name:
+                command += "SLOT %s " % quote_ident(slot_name, self)
+            # don't add "PHYSICAL", before 9.4 it was just START_REPLICATION XXX/XXX
+
+        else:
+            raise psycopg2.ProgrammingError("unrecognized replication type: %s" % repr(slot_type))
+
+        if type(start_lsn) is str:
+            lsn = start_lsn.split('/')
+            lsn = "%X/%08X" % (int(lsn[0], 16), int(lsn[1], 16))
+        else:
+            lsn = "%X/%08X" % ((start_lsn >> 32) & 0xFFFFFFFF, start_lsn & 0xFFFFFFFF)
+
+        command += lsn
+
+        if timeline != 0:
+            if slot_type == REPLICATION_LOGICAL:
+                raise psycopg2.ProgrammingError("cannot specify timeline for logical replication")
+
+            command += " TIMELINE %d" % timeline
+
+        if options:
+            if slot_type == REPLICATION_PHYSICAL:
+                raise psycopg2.ProgrammingError("cannot specify output plugin options for physical replication")
+
+            command += " ("
+            for k,v in options.iteritems():
+                if not command.endswith('('):
+                    command += ", "
+                command += "%s %s" % (quote_ident(k, self), _A(str(v)))
+            command += ")"
+
+        self.start_replication_expert(command, decode=decode)
+
+    # allows replication cursors to be used in select.select() directly
+    def fileno(self):
+        return self.connection.fileno()
+
+
 # a dbtype and adapter for Python UUID type
 
 class UUID_adapter(object):

psycopg/libpq_support.c

@@ -0,0 +1,104 @@
+/* libpq_support.c - functions not provided by libpq, but which are
+ * required for advanced communication with the server, such as
+ * streaming replication
+ *
+ * Copyright (C) 2003-2015 Federico Di Gregorio <fog@debian.org>
+ *
+ * This file is part of psycopg.
+ *
+ * psycopg2 is free software: you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link this program with the OpenSSL library (or with
+ * modified versions of OpenSSL that use the same license as OpenSSL),
+ * and distribute linked combinations including the two.
+ *
+ * You must obey the GNU Lesser General Public License in all respects for
+ * all of the code used other than OpenSSL.
+ *
+ * psycopg2 is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ */
+
+#define PSYCOPG_MODULE
+#include "psycopg/psycopg.h"
+
+#include "psycopg/libpq_support.h"
+
+/* htonl(), ntohl() */
+#ifdef _WIN32
+#include <winsock2.h>
+/* gettimeofday() */
+#include "psycopg/win32_support.h"
+#else
+#include <arpa/inet.h>
+#endif
+
+/* support routines taken from pg_basebackup/streamutil.c */
+
+/*
+ * Frontend version of GetCurrentTimestamp(), since we are not linked with
+ * backend code. The protocol always uses integer timestamps, regardless of
+ * server setting.
+ */
+pg_int64
+feGetCurrentTimestamp(void)
+{
+    pg_int64 result;
+    struct timeval tp;
+
+    gettimeofday(&tp, NULL);
+
+    result = (pg_int64) tp.tv_sec -
+        ((POSTGRES_EPOCH_JDATE - UNIX_EPOCH_JDATE) * SECS_PER_DAY);
+
+    result = (result * USECS_PER_SEC) + tp.tv_usec;
+
+    return result;
+}
+
+/*
+ * Converts an int64 to network byte order.
+ */
+void
+fe_sendint64(pg_int64 i, char *buf)
+{
+    uint32 n32;
+
+    /* High order half first, since we're doing MSB-first */
+    n32 = (uint32) (i >> 32);
+    n32 = htonl(n32);
+    memcpy(&buf[0], &n32, 4);
+
+    /* Now the low order half */
+    n32 = (uint32) i;
+    n32 = htonl(n32);
+    memcpy(&buf[4], &n32, 4);
+}
+
+/*
+ * Converts an int64 from network byte order to native format.
+ */
+pg_int64
+fe_recvint64(char *buf)
+{
+    pg_int64 result;
+    uint32 h32;
+    uint32 l32;
+
+    memcpy(&h32, buf, 4);
+    memcpy(&l32, buf + 4, 4);
+    h32 = ntohl(h32);
+    l32 = ntohl(l32);
+
+    result = h32;
+    result <<= 32;
+    result |= l32;
+
+    return result;
+}

psycopg/libpq_support.h

@@ -0,0 +1,49 @@
+/* libpq_support.h - definitions for libpq_support.c
+ *
+ * Copyright (C) 2003-2015 Federico Di Gregorio <fog@debian.org>
+ *
+ * This file is part of psycopg.
+ *
+ * psycopg2 is free software: you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link this program with the OpenSSL library (or with
+ * modified versions of OpenSSL that use the same license as OpenSSL),
+ * and distribute linked combinations including the two.
+ *
+ * You must obey the GNU Lesser General Public License in all respects for
+ * all of the code used other than OpenSSL.
+ *
+ * psycopg2 is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ */
+#ifndef PSYCOPG_LIBPQ_SUPPORT_H
+#define PSYCOPG_LIBPQ_SUPPORT_H 1
+
+#include "psycopg/config.h"
+#include "internal/c.h"
+
+/* type and constant definitions from internal postgres includes not available otherwise */
+typedef unsigned PG_INT64_TYPE XLogRecPtr;
+
+/* have to use lowercase %x, as PyString_FromFormat can't do %X */
+#define XLOGFMTSTR "%x/%x"
+#define XLOGFMTARGS(x) ((uint32)((x) >> 32)), ((uint32)((x) & 0xFFFFFFFF))
+
+/* Julian-date equivalents of Day 0 in Unix and Postgres reckoning */
+#define UNIX_EPOCH_JDATE	2440588 /* == date2j(1970, 1, 1) */
+#define POSTGRES_EPOCH_JDATE	2451545 /* == date2j(2000, 1, 1) */
+
+#define SECS_PER_DAY	86400
+#define USECS_PER_SEC	1000000LL
+
+HIDDEN pg_int64 feGetCurrentTimestamp(void);
+HIDDEN void fe_sendint64(pg_int64 i, char *buf);
+HIDDEN pg_int64 fe_recvint64(char *buf);
+
+#endif /* !defined(PSYCOPG_LIBPQ_SUPPORT_H) */