Implement Encode for Message and test Decode on ReadyForQuery

mason-postgres-protocol/src/decode.rs

@@ -4,7 +4,8 @@ use std::{io, str};
 
 pub trait Decode {
     fn decode(b: Bytes) -> io::Result<Self>
-    where Self: Sized;
+    where
+        Self: Sized;
 }
 
 #[inline]

mason-postgres-protocol/src/encode.rs

@@ -1,12 +1,13 @@
 use std::io;
 
 pub trait Encode {
-    fn size_hint() -> usize;
+    fn size_hint(&self) -> usize;
+
     fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()>;
 
     #[inline]
     fn to_bytes(&self) -> io::Result<Vec<u8>> {
-        let mut buf = Vec::with_capacity(Self::size_hint());
+        let mut buf = Vec::with_capacity(self.size_hint());
         self.encode(&mut buf)?;
         Ok(buf)
     }

mason-postgres-protocol/src/message.rs

@@ -1,17 +1,35 @@
-use crate::{Decode, NoticeResponse};
+use crate::{Decode, Encode, NoticeResponse, ReadyForQuery};
 use byteorder::{BigEndian, ReadBytesExt};
 use bytes::Bytes;
 use std::io::{self, Cursor};
 
 #[derive(Debug)]
-#[non_exhaustive]
 pub enum Message {
+    ReadyForQuery(ReadyForQuery),
     NoticeResponse(NoticeResponse),
 }
 
+impl Encode for Message {
+    fn size_hint(&self) -> usize {
+        match self {
+            Message::ReadyForQuery(body) => body.size_hint(),
+            Message::NoticeResponse(body) => body.size_hint(),
+        }
+    }
+
+    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
+        match self {
+            Message::ReadyForQuery(body) => body.encode(buf),
+            Message::NoticeResponse(body) => body.encode(buf),
+        }
+    }
+}
+
 impl Decode for Message {
     fn decode(b: Bytes) -> io::Result<Self>
-    where Self: Sized {
+    where
+        Self: Sized,
+    {
         let mut buf = Cursor::new(&b);
 
         let token = buf.read_u8()?;
@@ -22,7 +40,9 @@ impl Decode for Message {
         let b = b.slice(pos, pos + len - 4);
 
         Ok(match token {
+            // FIXME: These tokens are duplicated here and in the respective encode functions
             b'N' => Message::NoticeResponse(NoticeResponse::decode(b)?),
+            b'Z' => Message::ReadyForQuery(ReadyForQuery::decode(b)?),
 
             _ => unimplemented!("decode not implemented for token: {}", token as char),
         })

mason-postgres-protocol/src/notice_response.rs

@@ -1,7 +1,9 @@
-use crate::{decode::get_str, Decode};
+use crate::{decode::get_str, Decode, Encode};
+use byteorder::{BigEndian, WriteBytesExt};
 use bytes::Bytes;
 use std::{
-    fmt, io,
+    fmt,
+    io::{self, Write},
     pin::Pin,
     ptr::NonNull,
     str::{self, FromStr},
@@ -182,9 +184,26 @@ impl fmt::Debug for NoticeResponse {
     }
 }
 
+// FIXME: `Encode` here is (mostly) useless as its not easy to construct a NoticeResponse.
+//        Need a `NoticeResponse::builder().severity(...).build()` etc. type thing
+
+impl Encode for NoticeResponse {
+    fn size_hint(&self) -> usize { self.storage.len() + 5 }
+
+    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
+        buf.write_u8(b'Z')?;
+        buf.write_u32::<BigEndian>((4 + self.storage.len()) as u32)?;
+        buf.write_all(&self.storage)?;
+
+        Ok(())
+    }
+}
+
 impl Decode for NoticeResponse {
     fn decode(b: Bytes) -> io::Result<Self>
-    where Self: Sized {
+    where
+        Self: Sized,
+    {
         let storage = Pin::new(b);
 
         let mut code = None::<&str>;
@@ -384,7 +403,8 @@ mod tests {
 
             body
         })
-        .join().unwrap();
+        .join()
+        .unwrap();
 
         assert_eq!(body.code(), "42710");
         assert_eq!(body.routine(), Some("CreateExtension"));

mason-postgres-protocol/src/ready_for_query.rs

@@ -24,7 +24,7 @@ pub struct ReadyForQuery {
 
 impl Encode for ReadyForQuery {
     #[inline]
-    fn size_hint() -> usize { 6 }
+    fn size_hint(&self) -> usize { 6 }
 
     fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
         buf.write_u8(b'Z')?;
@@ -57,7 +57,8 @@ impl Decode for ReadyForQuery {
 #[cfg(test)]
 mod tests {
     use super::{ReadyForQuery, TransactionStatus};
-    use crate::{Decode, Encode};
+    use crate::{Decode, Encode, Message};
+    use bytes::Bytes;
     use std::io;
 
     #[test]
@@ -67,4 +68,22 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn it_decodes_ready_for_query() -> io::Result<()> {
+        // FIXME: A test-utils type thing could be useful here as these 7 lines are quite..
+        //        duplicated
+
+        let b = Bytes::from_static(b"Z\0\0\0\x05E");
+        let message = Message::decode(b)?;
+        let body = if let Message::ReadyForQuery(body) = message {
+            body
+        } else {
+            unreachable!();
+        };
+
+        assert_eq!(body.status, TransactionStatus::Error);
+
+        Ok(())
+    }
 }

rustfmt.toml

@@ -6,7 +6,6 @@ format_strings = true
 version = "Two"
 format_macro_matchers = true
 fn_single_line = true
-where_single_line = true
 reorder_impl_items = true
 condense_wildcard_suffixes = true
 use_field_init_shorthand = true