From a4b5cfc40871dfea22a0ebc6bc734876a5bd1e6a Mon Sep 17 00:00:00 2001
From: Zachery Gyurkovitz <zgyurkovitz@gmail.com>
Date: Fri, 19 Jul 2019 14:12:07 -0700
Subject: [PATCH] Implement `get_result`

---
 sqlx-postgres/src/connection/prepare.rs | 49 ++++++++++++++++++++++++-
 src/main.rs                             |  7 +++-
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/sqlx-postgres/src/connection/prepare.rs b/sqlx-postgres/src/connection/prepare.rs
index 336e682b..92d3289a 100644
--- a/sqlx-postgres/src/connection/prepare.rs
+++ b/sqlx-postgres/src/connection/prepare.rs
@@ -1,5 +1,5 @@
 use super::Connection;
-use sqlx_postgres_protocol::{self as proto, Execute, Message, Parse, Sync};
+use sqlx_postgres_protocol::{self as proto, DataRow, Execute, Message, Parse, Sync};
 use std::io;
 
 pub struct Prepare<'a> {
@@ -81,4 +81,49 @@ impl<'a> Prepare<'a> {
         // FIXME: This is an end-of-file error. How we should bubble this up here?
         unreachable!()
     }
-}
+
+    #[inline]
+    pub async fn get_result(self) -> io::Result<Option<DataRow>> {
+        proto::bind::trailer(
+            &mut self.connection.wbuf,
+            self.bind_state,
+            self.bind_values,
+            &[],
+        );
+
+        self.connection.send(Execute::new("", 0));
+        self.connection.send(Sync);
+        self.connection.flush().await?;
+
+        let mut row: Option<DataRow> = None;
+
+        while let Some(message) = self.connection.receive().await? {
+            match message {
+                Message::BindComplete | Message::ParseComplete => {
+                    // Indicates successful completion of a phase
+                }
+
+                Message::DataRow(data_row) => {
+                    // we only care about the first result.
+                    if row.is_none() {
+                        row = Some(data_row);
+                    }
+                }
+
+                Message::CommandComplete(_) => {}
+
+                Message::ReadyForQuery(_) => {
+                    // Successful completion of the whole cycle
+                    return Ok(row);
+                }
+
+                message => {
+                    unimplemented!("received {:?} unimplemented message", message);
+                }
+            }
+        }
+
+        // FIXME: This is an end-of-file error. How we should bubble this up here?
+        unreachable!()
+    }
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 59f9ecf6..a828905e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -31,11 +31,14 @@ CREATE TABLE IF NOT EXISTS users (
     .execute()
     .await?;
 
-    conn.prepare("INSERT INTO users (name) VALUES ($1)")
+    let row_id = conn
+        .prepare("INSERT INTO users (name) VALUES ($1) RETURNING id")
         .bind(b"Joe")
-        .execute()
+        .get_result()
         .await?;
 
+    println!("row_id: {:?}", row_id);
+
     let count = conn.prepare("SELECT name FROM users").execute().await?;
     println!("users: {}", count);