Implement IDENTIFY_SYSTEM command.

Author: jeff-davis

tokio-postgres/src/replication_client.rs

@@ -7,70 +7,136 @@ use bytes::Bytes;
 use fallible_iterator::FallibleIterator;
 use futures::{ready, Stream};
 use pin_project_lite::pin_project;
-use postgres_protocol::message::backend::{Field, Message, ReplicationMessage};
+use postgres_protocol::message::backend::{Message, ReplicationMessage};
 use std::io;
 use std::marker::PhantomPinned;
-use std::ops::Range;
 use std::pin::Pin;
+use std::str::from_utf8;
 use std::task::{Context, Poll};
 
+#[derive(Debug)]
+pub struct IdentifySystem {
+    systemid: String,
+    timeline: u32,
+    xlogpos: Lsn,
+    dbname: String,
+}
+
+impl IdentifySystem {
+    pub fn systemid(&self) -> &str {
+        &self.systemid
+    }
+
+    pub fn timeline(&self) -> u32 {
+        self.timeline
+    }
+
+    pub fn xlogpos(&self) -> Lsn {
+        self.xlogpos
+    }
+
+    pub fn dbname(&self) -> &str {
+        &self.dbname
+    }
+}
+
 /// Replication client connection.
 ///
 /// A replication client is used to issue replication commands, begin
 /// streaming, and send status updates to the server.
 pub struct ReplicationClient(pub(crate) Client);
 
 impl ReplicationClient {
+    /// IDENTIFY_SYSTEM message
+    pub async fn identify_system(&self) -> Result<IdentifySystem, Error> {
+        let iclient = self.0.inner();
+        let command = format!("IDENTIFY_SYSTEM");
+        let buf = simple_query::encode(iclient, &command)?;
+        let mut responses = iclient.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+        let rowdesc = match responses.next().await? {
+            Message::RowDescription(m) => m,
+            _ => return Err(Error::unexpected_message()),
+        };
+        let datarow = match responses.next().await? {
+            Message::DataRow(m) => m,
+            _ => return Err(Error::unexpected_message()),
+        };
+        match responses.next().await? {
+            Message::CommandComplete(_) => (),
+            _ => return Err(Error::unexpected_message()),
+        };
+        match responses.next().await? {
+            Message::ReadyForQuery(_) => (),
+            _ => return Err(Error::unexpected_message()),
+        };
+
+        let fields = rowdesc.fields().collect::<Vec<_>>().map_err(Error::parse)?;
+        let ranges = datarow.ranges().collect::<Vec<_>>().map_err(Error::parse)?;
+
+        if fields.len() != 4
+            || fields[0].type_oid() != Type::TEXT.oid()
+            || fields[1].type_oid() != Type::INT4.oid()
+            || fields[2].type_oid() != Type::TEXT.oid()
+            || fields[3].type_oid() != Type::TEXT.oid()
+            || ranges.len() != 4
+        {
+            return Err(Error::parse(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "expected (text, int4, text, text) result",
+            )));
+        };
+
+        let str_values = ranges
+            .iter()
+            .map(|r| from_utf8(&datarow.buffer()[r.to_owned().unwrap()]).unwrap())
+            .collect::<Vec<_>>();
+
+        Ok(IdentifySystem {
+            systemid: String::from(str_values[0]),
+            timeline: str_values[1].parse::<u32>().unwrap(),
+            xlogpos: Lsn::from(str_values[2]),
+            dbname: String::from(str_values[3]),
+        })
+    }
+
     /// show the value of the given setting
-    pub async fn show(&self, name: &str) -> Result<Vec<u8>, Error> {
+    pub async fn show(&self, name: &str) -> Result<String, Error> {
         let iclient = self.0.inner();
         let command = format!("SHOW \"{}\"", name);
         let buf = simple_query::encode(iclient, &command)?;
         let mut responses = iclient.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+        let rowdesc = match responses.next().await? {
+            Message::RowDescription(m) => m,
+            _ => return Err(Error::unexpected_message()),
+        };
+        let datarow = match responses.next().await? {
+            Message::DataRow(m) => m,
+            _ => return Err(Error::unexpected_message()),
+        };
         match responses.next().await? {
-            Message::RowDescription(rowdesc) => {
-                let fields: Vec<Field<'_>> = rowdesc.fields().collect().map_err(Error::parse)?;
-                if fields.len() != 1 {
-                    return Err(Error::parse(io::Error::new(
-                        io::ErrorKind::InvalidInput,
-                        "expected single column in response",
-                    )))
-                }
-                if fields[0].type_oid() != Type::TEXT.oid() {
-                    return Err(Error::parse(io::Error::new(
-                        io::ErrorKind::InvalidInput,
-                        "expected single text column in response",
-                    )))
-                }
-            }
+            Message::CommandComplete(_) => (),
             _ => return Err(Error::unexpected_message()),
         };
-
         match responses.next().await? {
-            Message::DataRow(d) => {
-                let ranges: Vec<Option<Range<usize>>> = d.ranges().collect().map_err(Error::parse)?;
-                if ranges.len() != 1 {
-                    return Err(Error::parse(io::Error::new(
-                        io::ErrorKind::InvalidInput,
-                        "expected single column in response",
-                    )))
-                }
-                // fetch only column
-                match &ranges[0] {
-                    Some(r) => Ok(Vec::from(&d.buffer()[r.to_owned()])),
-                    None => {
-                        Err(Error::parse(io::Error::new(
-                            io::ErrorKind::InvalidInput,
-                            "unexpected NULL setting in response",
-                        )))
-                    }
-                }
-            }
-            m => {
-                dbg!(m);
-                Err(Error::unexpected_message())
-            }
-        }
+            Message::ReadyForQuery(_) => (),
+            _ => return Err(Error::unexpected_message()),
+        };
+
+        let fields = rowdesc.fields().collect::<Vec<_>>().map_err(Error::parse)?;
+        let ranges = datarow.ranges().collect::<Vec<_>>().map_err(Error::parse)?;
+
+        if fields.len() != 1 || fields[0].type_oid() != Type::TEXT.oid() || ranges.len() != 1 {
+            return Err(Error::parse(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "expected single text column in response",
+            )));
+        };
+
+        let val = from_utf8(&datarow.buffer()[ranges[0].to_owned().unwrap()]).unwrap();
+
+        Ok(String::from(val))
     }
 
     /// Begin physical replication, consuming the replication client and producing a replication stream.

tokio-postgres/src/types.rs

@@ -6,6 +6,7 @@
 pub use postgres_types::*;
 
 /// Log Sequence Number for PostgreSQL Write-Ahead Log (transaction log).
+#[derive(Clone, Copy, Debug)]
 pub struct Lsn(u64);
 
 impl From<&str> for Lsn {