fmt code and use wal2json in tests

Author: ruslantalpa

docker/Dockerfile

@@ -1,3 +1,4 @@
 FROM postgres:12
 
 COPY sql_setup.sh /docker-entrypoint-initdb.d/
+RUN apt-get update && apt-get install postgresql-12-wal2json

docker/sql_setup.sh

@@ -66,6 +66,7 @@ ssl_cert_file = 'server.crt'
 ssl_key_file = 'server.key'
 wal_level = logical
 log_statement = 'all'
+shared_preload_libraries = 'wal2json'
 EOCONF
 
 cat > "$PGDATA/pg_hba.conf" <<-EOCONF

tokio-postgres/src/client.rs

@@ -2,8 +2,8 @@ use crate::codec::BackendMessages;
 use crate::config::{Host, SslMode};
 use crate::connection::{Request, RequestMessages};
 use crate::copy_out::CopyOutStream;
-use crate::replication::ReplicationStream;
 use crate::query::RowStream;
+use crate::replication::ReplicationStream;
 use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
@@ -12,8 +12,9 @@ use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    copy_in, copy_out, replication, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
-    Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+    copy_in, copy_out, prepare, query, replication, simple_query, slice_iter, CancelToken,
+    CopyInSink, Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
+    TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -435,14 +436,12 @@ impl Client {
     }
 
     /// TODO!
-    pub async fn start_replication(&self, query: &str) -> Result<ReplicationStream, Error>
-    {
+    pub async fn start_replication(&self, query: &str) -> Result<ReplicationStream, Error> {
         replication::start_replication(self.inner(), query).await
     }
 
     /// TODO!
-    pub async fn stop_replication(&self) -> Result<(), Error>
-    {
+    pub async fn stop_replication(&self) -> Result<(), Error> {
         replication::stop_replication(self.inner()).await
     }
 

tokio-postgres/src/config.rs

@@ -185,7 +185,7 @@ impl Config {
             keepalives_idle: Duration::from_secs(2 * 60 * 60),
             target_session_attrs: TargetSessionAttrs::Any,
             channel_binding: ChannelBinding::Prefer,
-            replication: None
+            replication: None,
         }
     }
 

tokio-postgres/src/lib.rs

@@ -121,12 +121,12 @@ pub use crate::config::Config;
 pub use crate::connection::Connection;
 pub use crate::copy_in::CopyInSink;
 pub use crate::copy_out::CopyOutStream;
-pub use crate::replication::ReplicationStream;
 use crate::error::DbError;
 pub use crate::error::Error;
 pub use crate::generic_client::GenericClient;
 pub use crate::portal::Portal;
 pub use crate::query::RowStream;
+pub use crate::replication::ReplicationStream;
 pub use crate::row::{Row, SimpleQueryRow};
 pub use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
@@ -158,13 +158,13 @@ mod connect_tls;
 mod connection;
 mod copy_in;
 mod copy_out;
-mod replication;
 pub mod error;
 mod generic_client;
 mod maybe_tls_stream;
 mod portal;
 mod prepare;
 mod query;
+mod replication;
 pub mod row;
 mod simple_query;
 #[cfg(feature = "runtime")]

tokio-postgres/src/replication.rs

@@ -4,36 +4,37 @@ use crate::connection::RequestMessages;
 use crate::{simple_query, Error};
 use bytes::{Bytes, BytesMut};
 use futures::{ready, Stream};
+use log::trace;
 use pin_project_lite::pin_project;
 use postgres_protocol::message::backend::Message;
-//use std::marker::PhantomPinned;
+use postgres_protocol::message::frontend;
+use std::marker::PhantomPinned;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use log::{trace};
-use postgres_protocol::message::frontend;
 
-pub async fn start_replication(client: &InnerClient, query: &str) -> Result<ReplicationStream, Error> {
+pub async fn start_replication(
+    client: &InnerClient,
+    query: &str,
+) -> Result<ReplicationStream, Error> {
     trace!("executing start replication query {}", query);
 
     let buf = simple_query::encode(client, query)?;
     let responses = start(client, buf).await?;
     Ok(ReplicationStream {
         responses,
-        // _p: PhantomPinned,
+        _p: PhantomPinned,
     })
 }
 
 pub async fn stop_replication(client: &InnerClient) -> Result<(), Error> {
     trace!("executing stop replication");
     let mut buf = BytesMut::new();
-        frontend::copy_done(&mut buf);
-        frontend::sync(&mut buf);
+    frontend::copy_done(&mut buf);
     let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf.freeze())))?;
     Ok(())
 }
 
 async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
-    
     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
     trace!("start in repication");
 
@@ -49,8 +50,8 @@ pin_project! {
     /// A stream of `START_REPLICATION` query data.
     pub struct ReplicationStream {
         responses: Responses,
-        // #[pin]
-        // _p: PhantomPinned,
+        #[pin]
+        _p: PhantomPinned,
     }
 }
 
@@ -64,10 +65,6 @@ impl Stream for ReplicationStream {
             Message::CopyData(body) => Poll::Ready(Some(Ok(body.into_bytes()))),
             Message::CopyDone => Poll::Ready(None),
             _ => Poll::Ready(Some(Err(Error::unexpected_message()))),
-            // m => {
-            //     debug!("msg: {:?}", m);
-            //     Poll::Ready(Some(Ok(Bytes::from("Hello world"))))
-            // }
         }
     }
 }

tokio-postgres/tests/test/main.rs

@@ -21,8 +21,8 @@ mod parse;
 #[cfg(feature = "runtime")]
 mod runtime;
 mod types;
-use std::thread;
 use std::sync::{Arc, Mutex};
+use std::thread;
 async fn connect_raw(s: &str) -> Result<(Client, Connection<TcpStream, NoTlsStream>), Error> {
     let socket = TcpStream::connect("127.0.0.1:5433").await.unwrap();
     let config = s.parse::<Config>().unwrap();
@@ -806,44 +806,72 @@ async fn replication_start() {
     let client = connect("user=postgres").await;
     let r_client = connect("user=postgres replication=database").await;
 
-    client.batch_execute(
-        "CREATE TABLE foo (
+    client
+        .batch_execute(
+            "CREATE TABLE IF NOT EXISTS foo (
             id SERIAL,
             name TEXT
         );
         DROP PUBLICATION IF EXISTS rust;
         CREATE PUBLICATION rust FOR ALL TABLES;
-        "
-    ).await.unwrap();
-    let _ = r_client.simple_query("DROP_REPLICATION_SLOT rust_slot").await;
-    let _ = r_client.simple_query("CREATE_REPLICATION_SLOT rust_slot LOGICAL pgoutput NOEXPORT_SNAPSHOT").await;
+        ",
+        )
+        .await
+        .unwrap();
+    let _ = r_client
+        .simple_query("DROP_REPLICATION_SLOT rust_slot")
+        .await;
+    // let _ = r_client.simple_query("CREATE_REPLICATION_SLOT rust_slot LOGICAL pgoutput NOEXPORT_SNAPSHOT").await;
+    let _ = r_client
+        .simple_query("CREATE_REPLICATION_SLOT rust_slot LOGICAL wal2json NOEXPORT_SNAPSHOT")
+        .await;
 
     let stream = r_client
-        .start_replication("START_REPLICATION SLOT rust_slot LOGICAL 0/0 (proto_version '1',  publication_names 'rust')")
+        //.start_replication("START_REPLICATION SLOT rust_slot LOGICAL 0/0 (proto_version '1',  publication_names 'rust')")
+        .start_replication("START_REPLICATION SLOT rust_slot LOGICAL 0/0 (\"pretty-print\" '1', \"format-version\" '2')")
         .await
         .unwrap();
-    
+
+    let events = Arc::new(Mutex::new(vec![]));
     let total_events = Arc::new(Mutex::new(0));
     let total_events_1 = Arc::clone(&total_events);
+    let events_1 = Arc::clone(&events);
     let t = tokio::spawn(async move {
-        let events = stream
-            .try_collect::<Vec<_>>()
-            .await
-            .unwrap();
+        let events = stream.try_collect::<Vec<_>>().await.unwrap();
         let mut num = total_events_1.lock().unwrap();
         *num = events.len();
+        let mut ev = events_1.lock().unwrap();
+        *ev = events;
     });
-    
-    client.batch_execute(
-        "
-            INSERT INTO foo (name) VALUES ('jim'), ('joe');
-            INSERT INTO foo (name) VALUES ('ann');
-            DROP TABLE foo;
-        "
-    ).await.unwrap();
+
+    client
+        .query("INSERT INTO foo (name) VALUES ('ann')", &[])
+        .await
+        .unwrap();
+    client
+        .query("INSERT INTO foo (name) VALUES ('jim'), ('joe')", &[])
+        .await
+        .unwrap();
+    client.query("DROP TABLE foo", &[]).await.unwrap();
+
     thread::sleep(time::Duration::from_secs(1)); //give a chance to pg to send the events
     let _ = r_client.stop_replication().await;
     let _ = t.await;
-    assert_eq!(*total_events.lock().unwrap(), 7);
+    println!("events {:?}", *events.lock().unwrap());
+    //assert_eq!(*total_events.lock().unwrap(), 2);
+    for e in &*events.lock().unwrap() {
+        match e[0].into() {
+            'k' => {
+                // let message_type = "keepalive";
+                // println!("type: ({}), {}", message_type, e.len());
+            }
+            'w' => {
+                let message_type = "dataframe";
+                let slice = e.slice(25..);
+                let data = std::str::from_utf8(slice.as_ref()).unwrap();
+                println!("type: ({}), {}, {}", message_type, e.len(), data);
+            }
+            _ => {}
+        };
+    }
 }
-