server version \d+ hack

src/client.rs

@@ -20,6 +20,7 @@ pub struct Client {
     buffer: BytesMut,
     name: String,
     cancel_mode: bool,
+    transaction_mode: bool,
     process_id: i32,
     secret_key: i32,
     client_server_map: ClientServerMap,
@@ -89,6 +90,7 @@ impl Client {
                         buffer: BytesMut::with_capacity(8196),
                         name: name,
                         cancel_mode: false,
+                        transaction_mode: true,
                         process_id: process_id,
                         secret_key: secret_key,
                         client_server_map: client_server_map,
@@ -108,6 +110,7 @@ impl Client {
                         buffer: BytesMut::with_capacity(8196),
                         name: String::from("cancel_mode"),
                         cancel_mode: true,
+                        transaction_mode: true,
                         process_id: process_id,
                         secret_key: secret_key,
                         client_server_map: client_server_map,
@@ -200,7 +203,7 @@ impl Client {
                         }
 
                         // Release server
-                        if !server.in_transaction() {
+                        if !server.in_transaction() && self.transaction_mode {
                             break;
                         }
                     }
@@ -253,7 +256,7 @@ impl Client {
                         }
 
                         // Release server
-                        if !server.in_transaction() {
+                        if !server.in_transaction() && self.transaction_mode {
                             break;
                         }
                     }
@@ -276,6 +279,12 @@ impl Client {
                                 return Err(err);
                             }
                         };
+
+                        // Release the server
+                        if !server.in_transaction() && self.transaction_mode {
+                            println!("Releasing after copy done");
+                            break;
+                        }
                     }
 
                     _ => {

src/messages.rs

@@ -6,6 +6,13 @@ use tokio::net::TcpStream;
 
 use crate::errors::Error;
 
+// This is a funny one. `psql` parses this to figure out which
+// queries to send when using shortcuts, e.g. \d+.
+//
+// TODO: Actually get the version from the server itself.
+//
+const SERVER_VESION: &str = "12.9 (Ubuntu 12.9-0ubuntu0.20.04.1)";
+
 pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
     let mut auth_ok = BytesMut::with_capacity(9);
 
@@ -18,6 +25,10 @@ pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
 
 pub async fn server_parameters(stream: &mut TcpStream) -> Result<(), Error> {
     let client_encoding = BytesMut::from(&b"client_encoding\0UTF8\0"[..]);
+    let server_version =
+        BytesMut::from(&format!("server_version\0{}\0", SERVER_VESION).as_bytes()[..]);
+
+    // Client encoding
     let len = client_encoding.len() as i32 + 4; // TODO: add more parameters here
     let mut res = BytesMut::with_capacity(len as usize + 1);
 
@@ -25,6 +36,11 @@ pub async fn server_parameters(stream: &mut TcpStream) -> Result<(), Error> {
     res.put_i32(len);
     res.put_slice(&client_encoding[..]);
 
+    let len = server_version.len() as i32 + 4;
+    res.put_u8(b'S');
+    res.put_i32(len);
+    res.put_slice(&server_version[..]);
+
     Ok(write_all(stream, res).await?)
 }