Merge pull request #2 from izik1/add-parameter-description-portal-suspended-describe

Add parameter description portal suspended describe

sqlx-postgres-protocol/src/bind.rs

@@ -1,6 +1,4 @@
-use crate::Encode;
 use byteorder::{BigEndian, ByteOrder};
-use std::io;
 
 // FIXME: Having structs here is breaking down. I think front-end messages should be
 //        simple functions that take the wbuf as a mut Vec

sqlx-postgres-protocol/src/describe.rs

@@ -0,0 +1,64 @@
+use crate::Encode;
+use std::io;
+
+#[derive(Debug)]
+pub enum DescribeKind {
+    Portal,
+    PreparedStatement,
+}
+
+#[derive(Debug)]
+pub struct Describe<'a> {
+    kind: DescribeKind,
+    name: &'a str,
+}
+
+impl<'a> Describe<'a> {
+    pub fn new(kind: DescribeKind, name: &'a str) -> Self {
+        Self { kind, name }
+    }
+}
+
+impl<'a> Encode for Describe<'a> {
+    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
+        buf.push(b'D');
+
+        let len = 4 + self.name.len() + 1 + 4;
+        buf.extend_from_slice(&(len as i32).to_be_bytes());
+
+        match &self.kind {
+            DescribeKind::Portal => buf.push(b'P'),
+            DescribeKind::PreparedStatement => buf.push(b'S'),
+        };
+
+        buf.extend_from_slice(self.name.as_bytes());
+        buf.push(b'\0');
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::{Describe, DescribeKind};
+    use crate::Encode;
+    use std::io;
+
+    #[test]
+    fn it_encodes_describe_portal() -> io::Result<()> {
+        let mut buf = vec![];
+        Describe::new(DescribeKind::Portal, "ABC123").encode(&mut buf)?;
+        assert_eq!(&buf, b"D\x00\x00\x00\x0fPABC123\x00");
+
+        Ok(())
+    }
+
+    #[test]
+    fn it_encodes_describe_statement() -> io::Result<()> {
+        let mut buf = vec![];
+        Describe::new(DescribeKind::PreparedStatement, "95 apples").encode(&mut buf)?;
+        assert_eq!(&buf, b"D\x00\x00\x00\x12S95 apples\x00");
+
+        Ok(())
+    }
+}

sqlx-postgres-protocol/src/lib.rs

@@ -6,10 +6,12 @@ pub mod bind;
 mod command_complete;
 mod data_row;
 mod decode;
+mod describe;
 mod encode;
 mod execute;
 mod message;
 mod notification_response;
+mod parameter_description;
 mod parameter_status;
 mod parse;
 mod password_message;
@@ -27,10 +29,12 @@ pub use self::{
     command_complete::CommandComplete,
     data_row::DataRow,
     decode::Decode,
+    describe::{Describe, DescribeKind},
     encode::Encode,
     execute::Execute,
     message::Message,
     notification_response::NotificationResponse,
+    parameter_description::ParameterDescription,
     parameter_status::ParameterStatus,
     parse::Parse,
     password_message::PasswordMessage,

sqlx-postgres-protocol/src/message.rs

@@ -1,6 +1,6 @@
 use crate::{
     Authentication, BackendKeyData, CommandComplete, DataRow, Decode, NotificationResponse,
-    ParameterStatus, ReadyForQuery, Response, RowDescription,
+    ParameterDescription, ParameterStatus, ReadyForQuery, Response, RowDescription,
 };
 use byteorder::{BigEndian, ByteOrder};
 use bytes::BytesMut;
@@ -20,6 +20,8 @@ pub enum Message {
     ParseComplete,
     BindComplete,
     NoData,
+    PortalSuspended,
+    ParameterDescription(ParameterDescription),
 }
 
 impl Message {
@@ -65,7 +67,8 @@ impl Message {
             b'1' => Message::ParseComplete,
             b'2' => Message::BindComplete,
             b'n' => Message::NoData,
-
+            b's' => Message::PortalSuspended,
+            b't' => Message::ParameterDescription(ParameterDescription::decode(src)?),
             _ => unimplemented!("decode not implemented for token: {}", token as char),
         };
 

sqlx-postgres-protocol/src/parameter_description.rs

@@ -0,0 +1,64 @@
+use crate::Decode;
+use byteorder::{BigEndian, ByteOrder};
+use bytes::Bytes;
+
+use std::io;
+
+type ObjectId = u32;
+
+#[derive(Debug)]
+pub struct ParameterDescription {
+    ids: Vec<ObjectId>,
+}
+
+impl Decode for ParameterDescription {
+    fn decode(src: Bytes) -> io::Result<Self> {
+        let count = BigEndian::read_u16(&*src) as usize;
+
+        // todo: error handling
+        assert_eq!(src.len(), count * 4 + 2);
+
+        let mut ids = Vec::with_capacity(count);
+        for i in 0..count {
+            let offset = i * 4 + 2; // 4==size_of(u32), 2==size_of(u16)
+            ids.push(BigEndian::read_u32(&src[offset..]));
+        }
+
+        Ok(ParameterDescription { ids })
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::ParameterDescription;
+    use crate::Decode;
+    use bytes::Bytes;
+    use std::io;
+
+    #[test]
+    fn it_decodes_parameter_description() -> io::Result<()> {
+        let src = Bytes::from_static(b"\x00\x02\x00\x00\x00\x00\x00\x00\x05\x00");
+        let desc = ParameterDescription::decode(src)?;
+
+        assert_eq!(desc.ids.len(), 2);
+        assert_eq!(desc.ids[0], 0x0000_0000);
+        assert_eq!(desc.ids[1], 0x0000_0500);
+        Ok(())
+    }
+
+    #[test]
+    fn it_decodes_empty_parameter_description() -> io::Result<()> {
+        let src = Bytes::from_static(b"\x00\x00");
+        let desc = ParameterDescription::decode(src)?;
+
+        assert_eq!(desc.ids.len(), 0);
+        Ok(())
+    }
+
+    #[test]
+    #[should_panic]
+    fn parameter_description_wrong_length_fails() -> () {
+        let src = Bytes::from_static(b"\x00\x00\x00\x01\x02\x03");
+        ParameterDescription::decode(src).unwrap();
+    }
+}

sqlx-postgres/src/connection/prepare.rs

@@ -16,7 +16,11 @@ pub fn prepare<'a, 'b>(connection: &'a mut Connection, query: &'b str) -> Prepar
 
     let bind_state = proto::bind::header(&mut connection.wbuf, "", "", &[]);
 
-    Prepare { connection, bind_state, bind_values: 0 }
+    Prepare {
+        connection,
+        bind_state,
+        bind_values: 0,
+    }
 }
 
 impl<'a> Prepare<'a> {
@@ -36,7 +40,12 @@ impl<'a> Prepare<'a> {
 
     #[inline]
     pub async fn execute(self) -> io::Result<u64> {
-        proto::bind::trailer(&mut self.connection.wbuf, self.bind_state, self.bind_values, &[]);
+        proto::bind::trailer(
+            &mut self.connection.wbuf,
+            self.bind_state,
+            self.bind_values,
+            &[],
+        );
 
         self.connection.send(Execute::new("", 0));
         self.connection.send(Sync);