docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-23 01:16:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

server version \d+ hack

Browse Source
This commit is contained in:
Lev Kokotov
2022-02-05 15:23:21 -08:00
parent 0e29d5e0ed
commit 7133d049c6
2 changed files with 27 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/client.rs
View File

@@ -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;
}
}
_ => {

16
src/messages.rs
View File

@@ -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?)
}