WIP Query

sqlx-postgres-protocol/src/lib.rs

@@ -7,6 +7,7 @@ mod encode;
 mod message;
 mod parameter_status;
 mod password_message;
+mod query;
 mod ready_for_query;
 mod response;
 mod startup_message;
@@ -20,6 +21,7 @@ pub use self::{
     message::Message,
     parameter_status::ParameterStatus,
     password_message::PasswordMessage,
+    query::Query,
     ready_for_query::{ReadyForQuery, TransactionStatus},
     response::{Response, ResponseBuilder, Severity},
     startup_message::StartupMessage,

sqlx-postgres-protocol/src/query.rs

@@ -0,0 +1,46 @@
+use crate::Encode;
+use bytes::BufMut;
+use std::io;
+
+#[derive(Debug)]
+pub struct Query<'a>(&'a str);
+
+impl<'a> Query<'a> {
+    pub fn new(query: &'a str) -> Self {
+        Self(query)
+    }
+}
+
+impl Encode for Query<'_> {
+    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
+        let len = self.0.len() + 4 + 1;
+        buf.reserve(len + 1);
+        buf.put_u8(b'Q');
+        buf.put_u32_be(len as u32);
+        buf.put(self.0);
+        buf.put_u8(0);
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Query;
+    use crate::Encode;
+    use std::io;
+
+    const QUERY_SELECT_1: &[u8] = b"Q\0\0\0\rSELECT 1\0";
+
+    #[test]
+    fn it_encodes_query() -> io::Result<()> {
+        let message = Query::new("SELECT 1");
+
+        let mut buf = Vec::new();
+        message.encode(&mut buf)?;
+
+        assert_eq!(&*buf, QUERY_SELECT_1);
+
+        Ok(())
+    }
+}

sqlx-postgres-protocol/src/terminate.rs

@@ -14,3 +14,24 @@ impl Encode for Terminate {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::Terminate;
+    use crate::Encode;
+    use std::io;
+
+    const TERMINATE: &[u8] = b"X\0\0\0\x04";
+
+    #[test]
+    fn it_encodes_terminate() -> io::Result<()> {
+        let message = Terminate;
+
+        let mut buf = Vec::new();
+        message.encode(&mut buf)?;
+
+        assert_eq!(&*buf, TERMINATE);
+
+        Ok(())
+    }
+}

sqlx-postgres/src/connection/establish.rs

@@ -38,7 +38,6 @@ pub async fn establish<'a, 'b: 'a>(
 
     conn.send(message).await?;
 
-    // FIXME: This feels like it could be reduced (see other connection flows)
     while let Some(message) = conn.stream.next().await {
         match message? {
             Message::Authentication(Authentication::Ok) => {

sqlx-postgres/src/connection/mod.rs

@@ -4,14 +4,13 @@ use futures::{
     task::{Context, Poll},
     Stream,
 };
-use std::fmt::Debug;
 use runtime::net::TcpStream;
 use sqlx_core::ConnectOptions;
 use sqlx_postgres_protocol::{Encode, Message, Terminate};
-use std::{io, pin::Pin};
+use std::{fmt::Debug, io, pin::Pin};
 
 mod establish;
-// mod query;
+mod query;
 
 pub struct Connection {
     stream: Framed<TcpStream>,
@@ -42,9 +41,9 @@ impl Connection {
         Ok(conn)
     }
 
-    // pub async fn execute<'a, 'b: 'a>(&'a mut self, query: &'b str) -> io::Result<()> {
-    //     query::query(self, query).await
-    // }
+    pub async fn execute<'a: 'b, 'b>(&'a mut self, query: &'b str) -> io::Result<()> {
+        query::query(self, query).await
+    }
 
     pub async fn close(mut self) -> io::Result<()> {
         self.send(Terminate).await?;

sqlx-postgres/src/connection/query.rs

@@ -1,22 +1,18 @@
 use super::Connection;
-use crate::protocol::{client::Query, server::Message as ServerMessage};
 use futures::StreamExt;
+use sqlx_postgres_protocol::{Message, Query};
 use std::io;
 
-pub async fn query<'a, 'b: 'a>(conn: &'a mut Connection, query: &'a str) -> io::Result<()> {
-    conn.send(Query { query }).await?;
+pub async fn query<'a: 'b, 'b>(conn: &'a mut Connection, query: &'b str) -> io::Result<()> {
+    conn.send(Query::new(query)).await?;
 
-    // FIXME: This feels like it could be reduced (see other connection flows)
-    while let Some(message) = conn.incoming.next().await {
-        match message {
-            ServerMessage::ReadyForQuery(_) => {
+    while let Some(message) = conn.stream.next().await {
+        match message? {
+            Message::ReadyForQuery(_) => {
                 break;
             }
 
-            ServerMessage::CommandComplete(body) => {
-            }
-
-            _ => {
+            message => {
                 unimplemented!("received {:?} unimplemented message", message);
             }
         }

src/main.rs

@@ -7,7 +7,7 @@ use std::io;
 async fn main() -> io::Result<()> {
     env_logger::init();
 
-    let conn = Connection::establish(
+    let mut conn = Connection::establish(
         ConnectOptions::new()
             .host("127.0.0.1")
             .port(5432)
@@ -16,6 +16,8 @@ async fn main() -> io::Result<()> {
     )
     .await?;
 
+    conn.execute("SELECT 1").await?;
+
     conn.close().await?;
 
     Ok(())