add blocking replication API

Add sync (blocking) wrappers for the async replication API.

Author: ericseppanen

Diff for: postgres/Cargo.toml

@@ -35,6 +35,7 @@ bytes = "1.0"
 fallible-iterator = "0.2"
 futures = "0.3"
 tokio-postgres = { version = "0.7.1", path = "../tokio-postgres" }
+postgres-protocol = { version = "0.6.1", path = "../postgres-protocol" }
 
 tokio = { version = "1.0", features = ["rt", "time"] }
 log = "0.4"

Diff for: postgres/src/client.rs

@@ -1,6 +1,6 @@
 use crate::connection::Connection;
 use crate::{
-    CancelToken, Config, CopyInWriter, CopyOutReader, Notifications, RowIter, Statement,
+    CancelToken, Config, CopyBoth, CopyInWriter, CopyOutReader, Notifications, RowIter, Statement,
     ToStatement, Transaction, TransactionBuilder,
 };
 use std::task::Poll;
@@ -395,6 +395,15 @@ impl Client {
         Ok(CopyOutReader::new(self.connection.as_ref(), stream))
     }
 
+    /// Executes a CopyBoth query, returning a combined Stream+Sink type to read and write copy
+    /// data.
+    pub fn copy_both_simple(&mut self, query: &str) -> Result<CopyBoth<'_>, Error> {
+        let stream = self
+            .connection
+            .block_on(self.client.copy_both_simple(query))?;
+        Ok(CopyBoth::new(self.connection.as_ref(), stream))
+    }
+
     /// Executes a sequence of SQL statements using the simple query protocol.
     ///
     /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that

Diff for: postgres/src/copy_both.rs

@@ -0,0 +1,77 @@
+use crate::connection::ConnectionRef;
+use crate::lazy_pin::LazyPin;
+use bytes::{Buf, Bytes, BytesMut};
+use futures::{SinkExt, StreamExt};
+use std::io::{self, BufRead, Read};
+use tokio_postgres::{CopyBothDuplex, Error};
+
+/// The reader/writer returned by the `copy_both_simple` method.
+pub struct CopyBoth<'a> {
+    pub(crate) connection: ConnectionRef<'a>,
+    pub(crate) stream_sink: LazyPin<CopyBothDuplex<Bytes>>,
+    buf: BytesMut,
+    cur: Bytes,
+}
+
+impl<'a> CopyBoth<'a> {
+    pub(crate) fn new(
+        connection: ConnectionRef<'a>,
+        duplex: CopyBothDuplex<Bytes>,
+    ) -> CopyBoth<'a> {
+        CopyBoth {
+            connection,
+            stream_sink: LazyPin::new(duplex),
+            buf: BytesMut::new(),
+            cur: Bytes::new(),
+        }
+    }
+
+    /// Completes the copy, returning the number of rows written.
+    ///
+    /// If this is not called, the copy will be aborted.
+    pub fn finish(mut self) -> Result<u64, Error> {
+        self.flush_inner()?;
+        self.connection.block_on(self.stream_sink.pinned().finish())
+    }
+
+    fn flush_inner(&mut self) -> Result<(), Error> {
+        if self.buf.is_empty() {
+            return Ok(());
+        }
+
+        self.connection
+            .block_on(self.stream_sink.pinned().send(self.buf.split().freeze()))
+    }
+}
+
+impl Read for CopyBoth<'_> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let b = self.fill_buf()?;
+        let len = usize::min(buf.len(), b.len());
+        buf[..len].copy_from_slice(&b[..len]);
+        self.consume(len);
+        Ok(len)
+    }
+}
+
+impl BufRead for CopyBoth<'_> {
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        while !self.cur.has_remaining() {
+            let mut stream = self.stream_sink.pinned();
+            match self
+                .connection
+                .block_on(async { stream.next().await.transpose() })
+            {
+                Ok(Some(cur)) => self.cur = cur,
+                Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
+                Ok(None) => break,
+            };
+        }
+
+        Ok(&self.cur)
+    }
+
+    fn consume(&mut self, amt: usize) {
+        self.cur.advance(amt);
+    }
+}

Diff for: postgres/src/lib.rs

@@ -72,6 +72,7 @@ pub use tokio_postgres::{
 pub use crate::cancel_token::CancelToken;
 pub use crate::client::*;
 pub use crate::config::Config;
+pub use crate::copy_both::CopyBoth;
 pub use crate::copy_in_writer::CopyInWriter;
 pub use crate::copy_out_reader::CopyOutReader;
 #[doc(no_inline)]
@@ -92,14 +93,18 @@ mod cancel_token;
 mod client;
 pub mod config;
 mod connection;
+mod copy_both;
 mod copy_in_writer;
 mod copy_out_reader;
 mod generic_client;
 mod lazy_pin;
 pub mod notifications;
+pub mod replication;
 mod row_iter;
 mod transaction;
 mod transaction_builder;
 
+#[cfg(test)]
+mod replication_test;
 #[cfg(test)]
 mod test;

Diff for: postgres/src/replication.rs

@@ -0,0 +1,164 @@
+//! Utilities for working with the PostgreSQL replication copy both format.
+
+use crate::connection::ConnectionRef;
+use crate::lazy_pin::LazyPin;
+use crate::{CopyBoth, Error};
+use bytes::Bytes;
+use fallible_iterator::FallibleIterator;
+use futures::Stream;
+use postgres_protocol::message::backend::{LogicalReplicationMessage, ReplicationMessage};
+use std::task::Poll;
+use std::time::SystemTime;
+use tokio_postgres::replication::{LogicalReplicationStream, ReplicationStream};
+use tokio_postgres::types::PgLsn;
+
+/// A type which deserializes the postgres replication protocol.
+///
+/// This type can be used with both physical and logical replication to get
+/// access to the byte content of each replication message.
+///
+/// This is the sync (blocking) version of [`ReplicationStream`]
+pub struct ReplicationIter<'a> {
+    connection: ConnectionRef<'a>,
+    stream: LazyPin<ReplicationStream>,
+}
+
+impl<'a> ReplicationIter<'a> {
+    /// Creates a new `ReplicationIter`.
+    pub fn new(copyboth: CopyBoth<'a>) -> Self {
+        let unpinned_copyboth = copyboth
+            .stream_sink
+            .into_unpinned()
+            .expect("copy-both stream has already been used");
+        let stream = ReplicationStream::new(unpinned_copyboth);
+        Self {
+            connection: copyboth.connection,
+            stream: LazyPin::new(stream),
+        }
+    }
+
+    /// Send standby update to server.
+    pub fn standby_status_update(
+        &mut self,
+        write_lsn: PgLsn,
+        flush_lsn: PgLsn,
+        apply_lsn: PgLsn,
+        timestamp: SystemTime,
+        reply: u8,
+    ) -> Result<(), Error> {
+        self.connection.block_on(
+            self.stream
+                .pinned()
+                .standby_status_update(write_lsn, flush_lsn, apply_lsn, timestamp, reply),
+        )
+    }
+
+    /// Send hot standby feedback message to server.
+    pub fn hot_standby_feedback(
+        &mut self,
+        timestamp: SystemTime,
+        global_xmin: u32,
+        global_xmin_epoch: u32,
+        catalog_xmin: u32,
+        catalog_xmin_epoch: u32,
+    ) -> Result<(), Error> {
+        self.connection
+            .block_on(self.stream.pinned().hot_standby_feedback(
+                timestamp,
+                global_xmin,
+                global_xmin_epoch,
+                catalog_xmin,
+                catalog_xmin_epoch,
+            ))
+    }
+}
+
+impl<'a> FallibleIterator for ReplicationIter<'a> {
+    type Item = ReplicationMessage<Bytes>;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        let pinstream = &mut self.stream;
+
+        self.connection
+            .poll_block_on(|cx, _, _| match pinstream.pinned().poll_next(cx) {
+                Poll::Ready(x) => Poll::Ready(x.transpose()),
+                Poll::Pending => Poll::Pending,
+            })
+    }
+}
+
+/// A type which deserializes the postgres logical replication protocol. This
+/// type gives access to a high level representation of the changes in
+/// transaction commit order.
+///
+/// This is the sync (blocking) version of [`LogicalReplicationStream`]
+pub struct LogicalReplicationIter<'a> {
+    connection: ConnectionRef<'a>,
+    stream: LazyPin<LogicalReplicationStream>,
+}
+
+impl<'a> LogicalReplicationIter<'a> {
+    /// Creates a new `ReplicationThing`.
+    pub fn new(copyboth: CopyBoth<'a>) -> Self {
+        let unpinned_copyboth = copyboth
+            .stream_sink
+            .into_unpinned()
+            .expect("copy-both stream has already been used");
+        let stream = LogicalReplicationStream::new(unpinned_copyboth);
+        Self {
+            connection: copyboth.connection,
+            stream: LazyPin::new(stream),
+        }
+    }
+
+    /// Send standby update to server.
+    pub fn standby_status_update(
+        &mut self,
+        write_lsn: PgLsn,
+        flush_lsn: PgLsn,
+        apply_lsn: PgLsn,
+        timestamp: SystemTime,
+        reply: u8,
+    ) -> Result<(), Error> {
+        self.connection.block_on(
+            self.stream
+                .pinned()
+                .standby_status_update(write_lsn, flush_lsn, apply_lsn, timestamp, reply),
+        )
+    }
+
+    /// Send hot standby feedback message to server.
+    pub fn hot_standby_feedback(
+        &mut self,
+        timestamp: SystemTime,
+        global_xmin: u32,
+        global_xmin_epoch: u32,
+        catalog_xmin: u32,
+        catalog_xmin_epoch: u32,
+    ) -> Result<(), Error> {
+        self.connection
+            .block_on(self.stream.pinned().hot_standby_feedback(
+                timestamp,
+                global_xmin,
+                global_xmin_epoch,
+                catalog_xmin,
+                catalog_xmin_epoch,
+            ))
+    }
+}
+
+impl<'a> FallibleIterator for LogicalReplicationIter<'a> {
+    type Item = ReplicationMessage<LogicalReplicationMessage>;
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        let pinstream = &mut self.stream;
+
+        self.connection
+            .poll_block_on(|cx, _, _| match pinstream.pinned().poll_next(cx) {
+                Poll::Ready(x) => Poll::Ready(x.transpose()),
+                Poll::Pending => Poll::Pending,
+            })
+    }
+}