From 381f06d46f190bd2812016d9660176a6ed95bcb5 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <lev.kokotov@gmail.com>
Date: Fri, 4 Feb 2022 08:26:50 -0800
Subject: [PATCH] COPY OUT support

---
 src/client.rs | 41 +++++++++++++++++++++++++++--------------
 src/server.rs | 21 +++++++++++++++++++++
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 9298665..d530a4b 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -118,14 +118,21 @@ impl Client {
                 match code {
                     'Q' => {
                         server.send(original).await?;
-                        let response = server.recv().await?;
-                        match write_all_half(&mut self.write, response).await {
-                            Ok(_) => (),
-                            Err(err) => {
-                                server.mark_bad();
-                                return Err(err);
+
+                        loop {
+                            let response = server.recv().await?;
+                            match write_all_half(&mut self.write, response).await {
+                                Ok(_) => (),
+                                Err(err) => {
+                                    server.mark_bad();
+                                    return Err(err);
+                                }
+                            };
+
+                            if !server.is_data_available() {
+                                break;
                             }
-                        };
+                        }
 
                         // Release server
                         if !server.in_transaction() {
@@ -166,14 +173,20 @@ impl Client {
                         server.send(self.buffer.clone()).await?;
                         self.buffer.clear();
 
-                        let response = server.recv().await?;
-                        match write_all_half(&mut self.write, response).await {
-                            Ok(_) => (),
-                            Err(err) => {
-                                server.mark_bad();
-                                return Err(err);
+                        loop {
+                            let response = server.recv().await?;
+                            match write_all_half(&mut self.write, response).await {
+                                Ok(_) => (),
+                                Err(err) => {
+                                    server.mark_bad();
+                                    return Err(err);
+                                }
+                            };
+
+                            if !server.is_data_available() {
+                                break;
                             }
-                        };
+                        }
 
                         // Release server
                         if !server.in_transaction() {
diff --git a/src/server.rs b/src/server.rs
index 098a77e..b83bbb7 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -17,6 +17,7 @@ pub struct Server {
     backend_id: i32,
     secret_key: i32,
     in_transaction: bool,
+    data_available: bool,
     bad: bool,
 }
 
@@ -136,6 +137,7 @@ impl Server {
                         backend_id: backend_id,
                         secret_key: secret_key,
                         in_transaction: false,
+                        data_available: false,
                         bad: false,
                     });
                 }
@@ -208,6 +210,21 @@ impl Server {
                 // CopyInResponse: copy is starting from client to server
                 'G' => break,
 
+                // CopyOutResponse: copy is starting from the server to the client
+                'H' => {
+                    self.data_available = true;
+                    break;
+                }
+
+                // CopyData
+                'd' => break,
+
+                // CopyDone
+                'c' => {
+                    self.data_available = false;
+                    // Buffer until ReadyForQuery shows up
+                }
+
                 _ => {
                     // Keep buffering,
                 }
@@ -224,6 +241,10 @@ impl Server {
         self.in_transaction
     }
 
+    pub fn is_data_available(&self) -> bool {
+        self.data_available
+    }
+
     pub fn is_bad(&self) -> bool {
         self.bad
     }