From 7133d049c6c709ae43eddc947f1e852231820541 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <lev.kokotov@gmail.com>
Date: Sat, 5 Feb 2022 15:23:21 -0800
Subject: [PATCH] server version \d+ hack

---
 src/client.rs   | 13 +++++++++++--
 src/messages.rs | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 0fe1859..4accc8e 100644
--- a/src/client.rs
+++ b/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;
+                        }
                     }
 
                     _ => {
diff --git a/src/messages.rs b/src/messages.rs
index 7c07669..3f91487 100644
--- a/src/messages.rs
+++ b/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?)
 }