feat(postgres): add more frontend protocol messages

sqlx-postgres/Cargo.toml

@@ -41,6 +41,7 @@ memchr = "2.3"
 bitflags = "1.2"
 base64 = "0.13.0"
 md-5 = "0.9.1"
+itoa = "0.4.7"
 
 [dev-dependencies]
 sqlx-core = { version = "0.6.0-pre", path = "../sqlx-core", features = ["_mock"] }

sqlx-postgres/src/protocol/frontend.rs

@@ -1,7 +1,29 @@
+mod bind;
+mod close;
+mod describe;
+mod execute;
+mod flush;
+mod parse;
 mod password;
+mod portal;
+mod query;
 mod startup;
+mod statement;
+mod sync;
+mod target;
 mod terminate;
 
+pub(crate) use bind::Bind;
+pub(crate) use close::Close;
+pub(crate) use describe::Describe;
+pub(crate) use execute::Execute;
+pub(crate) use flush::Flush;
+pub(crate) use parse::Parse;
 pub(crate) use password::{Password, PasswordMd5};
+pub(crate) use portal::PortalRef;
+pub(crate) use query::Query;
 pub(crate) use startup::Startup;
+pub(crate) use statement::StatementRef;
+pub(crate) use sync::Sync;
+pub(crate) use target::Target;
 pub(crate) use terminate::Terminate;

sqlx-postgres/src/protocol/frontend/bind.rs

@@ -0,0 +1,41 @@
+use crate::io::PgWriteExt;
+use crate::protocol::frontend::{PortalRef, StatementRef};
+use crate::PgArguments;
+use sqlx_core::io::{Serialize, WriteExt};
+use sqlx_core::Result;
+
+pub(crate) struct Bind<'a> {
+    pub(crate) portal: PortalRef,
+    pub(crate) statement: StatementRef,
+    pub(crate) arguments: &'a PgArguments<'a>,
+}
+
+impl Serialize<'_> for Bind<'_> {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'B');
+        buf.write_len_prefixed(|buf| {
+            self.portal.serialize(buf)?;
+            self.statement.serialize(buf)?;
+
+            // the parameter format codes, each must presently be zero (text) or one (binary)
+            // can use one to indicate that all parameters use that format
+            write_i16_arr(buf, &[1]);
+
+            todo!("arguments");
+
+            // the result format codes, each must presently be zero (text) or one (binary)
+            // can use one to indicate that all results use that format
+            write_i16_arr(buf, &[1]);
+
+            Ok(())
+        })
+    }
+}
+
+fn write_i16_arr(buf: &mut Vec<u8>, arr: &[i16]) {
+    buf.extend(&(arr.len() as i16).to_be_bytes());
+
+    for val in arr {
+        buf.extend(&val.to_be_bytes());
+    }
+}

sqlx-postgres/src/protocol/frontend/close.rs

@@ -0,0 +1,16 @@
+use crate::io::PgWriteExt;
+use crate::protocol::frontend::Target;
+use sqlx_core::io::{Serialize, WriteExt};
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Close {
+    target: Target,
+}
+
+impl Serialize<'_> for Close {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'C');
+        buf.write_len_prefixed(|buf| self.target.serialize(buf))
+    }
+}

sqlx-postgres/src/protocol/frontend/describe.rs

@@ -0,0 +1,16 @@
+use crate::io::PgWriteExt;
+use crate::protocol::frontend::Target;
+use sqlx_core::io::{Serialize, WriteExt};
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Describe {
+    target: Target,
+}
+
+impl Serialize<'_> for Describe {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'D');
+        buf.write_len_prefixed(|buf| self.target.serialize(buf))
+    }
+}

sqlx-postgres/src/protocol/frontend/execute.rs

@@ -0,0 +1,25 @@
+use crate::io::PgWriteExt;
+use crate::protocol::frontend::PortalRef;
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Execute {
+    pub(crate) portal: PortalRef,
+
+    /// Maximum number of rows to return, if portal contains a query
+    /// that returns rows (ignored otherwise). Zero denotes “no limit”.
+    pub(crate) max_rows: u32,
+}
+
+impl Serialize<'_> for Execute {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'E');
+        buf.write_len_prefixed(|buf| {
+            self.portal.serialize(buf)?;
+            buf.extend(&self.max_rows.to_be_bytes());
+
+            Ok(())
+        })
+    }
+}

sqlx-postgres/src/protocol/frontend/flush.rs

@@ -0,0 +1,13 @@
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Flush;
+
+impl Serialize<'_> for Flush {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'H');
+
+        Ok(())
+    }
+}

sqlx-postgres/src/protocol/frontend/parse.rs

@@ -0,0 +1,37 @@
+use crate::io::PgWriteExt;
+use crate::protocol::frontend::{PortalRef, StatementRef};
+use sqlx_core::io::{Serialize, WriteExt};
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Parse<'a> {
+    pub(crate) statement: StatementRef,
+    pub(crate) sql: &'a str,
+
+    /// The parameter data types specified (could be zero). Note that this is not an
+    /// indication of the number of parameters that might appear in the query string,
+    /// only the number that the frontend wants to pre-specify types for.
+    pub(crate) parameters: &'a [u32],
+}
+
+impl Serialize<'_> for Parse<'_> {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'P');
+        buf.write_len_prefixed(|buf| {
+            self.statement.serialize(buf)?;
+
+            buf.write_str_nul(self.sql);
+
+            // TODO: return a proper error
+            assert!(!(self.parameters.len() >= (u16::MAX as usize)));
+
+            buf.extend(&(self.parameters.len() as u16).to_be_bytes());
+
+            for &oid in self.parameters {
+                buf.extend(&oid.to_be_bytes());
+            }
+
+            Ok(())
+        })
+    }
+}

sqlx-postgres/src/protocol/frontend/portal.rs

@@ -0,0 +1,22 @@
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) enum PortalRef {
+    Unnamed,
+    Named(u32),
+}
+
+impl Serialize<'_> for PortalRef {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        if let PortalRef::Named(id) = self {
+            buf.extend_from_slice(b"_sqlx_p_");
+
+            itoa::write(&mut *buf, *id).unwrap();
+        }
+
+        buf.push(b'\0');
+
+        Ok(())
+    }
+}

sqlx-postgres/src/protocol/frontend/query.rs

@@ -0,0 +1,21 @@
+use crate::io::PgWriteExt;
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Query<'a>(pub(crate) &'a str);
+
+impl Serialize<'_> for Query<'_> {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.reserve(1 + self.0.len() + 1 + 4);
+
+        buf.push(b'Q');
+
+        buf.write_len_prefixed(|buf| {
+            buf.extend_from_slice(self.0.as_bytes());
+            buf.push(0);
+
+            Ok(())
+        })
+    }
+}

sqlx-postgres/src/protocol/frontend/statement.rs

@@ -0,0 +1,22 @@
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) enum StatementRef {
+    Unnamed,
+    Named(u32),
+}
+
+impl Serialize<'_> for StatementRef {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        if let StatementRef::Named(id) = self {
+            buf.extend_from_slice(b"_sqlx_s_");
+
+            itoa::write(&mut *buf, *id).unwrap();
+        }
+
+        buf.push(b'\0');
+
+        Ok(())
+    }
+}

sqlx-postgres/src/protocol/frontend/sync.rs

@@ -0,0 +1,13 @@
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+#[derive(Debug)]
+pub(crate) struct Sync;
+
+impl Serialize<'_> for Sync {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.push(b'S');
+
+        Ok(())
+    }
+}

sqlx-postgres/src/protocol/frontend/target.rs

@@ -0,0 +1,32 @@
+use crate::io::PgWriteExt;
+use crate::protocol::frontend::{PortalRef, StatementRef};
+use sqlx_core::io::Serialize;
+use sqlx_core::Result;
+
+/// Target a command at a portal *or* statement.
+/// Used by [`Describe`] and [`Close`].
+#[derive(Debug)]
+pub(crate) enum Target {
+    Portal(PortalRef),
+    Statement(StatementRef),
+}
+
+impl Serialize<'_> for Target {
+    fn serialize_with(&self, buf: &mut Vec<u8>, _: ()) -> Result<()> {
+        buf.write_len_prefixed(|buf| {
+            match self {
+                Self::Portal(portal) => {
+                    buf.push(b'P');
+                    portal.serialize(buf);
+                }
+
+                Self::Statement(statement) => {
+                    buf.push(b'S');
+                    statement.serialize(buf);
+                }
+            }
+
+            Ok(())
+        })
+    }
+}