reqs

src/server.rs

@@ -38,6 +38,7 @@ pub struct Server {
 
     /// Our server response buffer. We buffer data before we give it to the client.
     buffer: BytesMut,
+    is_async: bool,
 
     /// Server information the server sent us over on startup.
     server_info: BytesMut,
@@ -450,6 +451,7 @@ impl Server {
                         read: BufReader::new(read),
                         write,
                         buffer: BytesMut::with_capacity(8196),
+                        is_async: false,
                         server_info,
                         process_id,
                         secret_key,
@@ -537,6 +539,16 @@ impl Server {
         }
     }
 
+    /// Switch to async mode, flushing messages as soon
+    /// as we receive them without buffering or waiting for "ReadyForQuery".
+    pub fn switch_async(&mut self, on: bool) {
+        if on {
+            self.is_async = true;
+        } else {
+            self.is_async = false;
+        }
+    }
+
     /// Receive data from the server in response to a client request.
     /// This method must be called multiple times while `self.is_data_available()` is true
     /// in order to receive all data the server has to offer.
@@ -557,8 +569,6 @@ impl Server {
             let code = message.get_u8() as char;
             let _len = message.get_i32();
 
-            trace!("Message: {}", code);
-
             match code {
                 // ReadyForQuery
                 'Z' => {
@@ -632,7 +642,10 @@ impl Server {
                 // DataRow
                 'D' => {
                     // More data is available after this message, this is not the end of the reply.
-                    self.data_available = true;
+                    // If we're async, flush to client now.
+                    if !self.is_async {
+                        self.data_available = true;
+                    }
 
                     // Don't flush yet, the more we buffer, the faster this goes...up to a limit.
                     if self.buffer.len() >= 8196 {
@@ -645,7 +658,10 @@ impl Server {
 
                 // CopyOutResponse: copy is starting from the server to the client.
                 'H' => {
-                    self.data_available = true;
+                    // If we're in async mode, flush now.
+                    if !self.is_async {
+                        self.data_available = true;
+                    }
                     break;
                 }
 
@@ -665,6 +681,10 @@ impl Server {
                 // Keep buffering until ReadyForQuery shows up.
                 _ => (),
             };
+
+            if self.is_async {
+                break;
+            }
         }
 
         let bytes = self.buffer.clone();