docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-25 18:06:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make infer role configurable and fix double parse bug (#533)

Browse Source 
* Make infer role configurable and fix double parse bug

* Fix tests

* Enable infer_role_from query in toml for tests

* Fix test

* Add max length config, add logging for which application is failing to parse, and change config name

* fmt

* Update src/config.rs

---------

Co-authored-by: Lev Kokotov <levkk@users.noreply.github.com>
This commit is contained in:
Zain Kabani
2023-08-08 16:10:03 -04:00
committed by GitHub
parent 7c3c90c38e
commit e14b283f0c
8 changed files with 197 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
src/client.rs
View File

@@ -774,6 +774,9 @@ where
let mut prepared_statement = None;
let mut will_prepare = false;
let client_identifier =
ClientIdentifier::new(&self.application_name, &self.username, &self.pool_name);
// Our custom protocol loop.
// We expect the client to either start a transaction with regular queries
// or issue commands for our sharding and server selection protocol.
@@ -812,6 +815,21 @@ where
message_result = read_message(&mut self.read) => message_result?
};
// Handle admin database queries.
if self.admin {
debug!("Handling admin command");
handle_admin(&mut self.write, message, self.client_server_map.clone()).await?;
continue;
}
// Get a pool instance referenced by the most up-to-date
// pointer. This ensures we always read the latest config
// when starting a query.
let mut pool = self.get_pool().await?;
query_router.update_pool_settings(pool.settings.clone());
let mut initial_parsed_ast = None;
match message[0] as char {
// Buffer extended protocol messages even if we do not have
// a server connection yet. Hopefully, when we get the S message
@@ -841,24 +859,34 @@ where
'Q' => {
if query_router.query_parser_enabled() {
if let Ok(ast) = QueryRouter::parse(&message) {
let plugin_result = query_router.execute_plugins(&ast).await;
match query_router.parse(&message) {
Ok(ast) => {
let plugin_result = query_router.execute_plugins(&ast).await;
match plugin_result {
Ok(PluginOutput::Deny(error)) => {
error_response(&mut self.write, &error).await?;
continue;
}
match plugin_result {
Ok(PluginOutput::Deny(error)) => {
error_response(&mut self.write, &error).await?;
continue;
}
Ok(PluginOutput::Intercept(result)) => {
write_all(&mut self.write, result).await?;
continue;
}
Ok(PluginOutput::Intercept(result)) => {
write_all(&mut self.write, result).await?;
continue;
}
_ => (),
};
_ => (),
};
let _ = query_router.infer(&ast);
let _ = query_router.infer(&ast);
initial_parsed_ast = Some(ast);
}
Err(error) => {
warn!(
"Query parsing error: {} (client: {})",
error, client_identifier
);
}
}
}
}
@@ -872,13 +900,21 @@ where
self.buffer.put(&message[..]);
if query_router.query_parser_enabled() {
if let Ok(ast) = QueryRouter::parse(&message) {
if let Ok(output) = query_router.execute_plugins(&ast).await {
plugin_output = Some(output);
}
match query_router.parse(&message) {
Ok(ast) => {
if let Ok(output) = query_router.execute_plugins(&ast).await {
plugin_output = Some(output);
}
let _ = query_router.infer(&ast);
}
let _ = query_router.infer(&ast);
}
Err(error) => {
warn!(
"Query parsing error: {} (client: {})",
error, client_identifier
);
}
};
}
continue;
@@ -922,13 +958,6 @@ where
_ => (),
}
// Handle admin database queries.
if self.admin {
debug!("Handling admin command");
handle_admin(&mut self.write, message, self.client_server_map.clone()).await?;
continue;
}
// Check on plugin results.
match plugin_output {
Some(PluginOutput::Deny(error)) => {
@@ -941,11 +970,6 @@ where
_ => (),
};
// Get a pool instance referenced by the most up-to-date
// pointer. This ensures we always read the latest config
// when starting a query.
let mut pool = self.get_pool().await?;
// Check if the pool is paused and wait until it's resumed.
if pool.wait_paused().await {
// Refresh pool information, something might have changed.
@@ -1165,6 +1189,9 @@ where
None => {
trace!("Waiting for message inside transaction or in session mode");
// This is not an initial message so discard the initial_parsed_ast
initial_parsed_ast.take();
match tokio::time::timeout(
idle_client_timeout_duration,
read_message(&mut self.read),
@@ -1221,7 +1248,22 @@ where
// Query
'Q' => {
if query_router.query_parser_enabled() {
if let Ok(ast) = QueryRouter::parse(&message) {
// We don't want to parse again if we already parsed it as the initial message
let ast = match initial_parsed_ast {
Some(_) => Some(initial_parsed_ast.take().unwrap()),
None => match query_router.parse(&message) {
Ok(ast) => Some(ast),
Err(error) => {
warn!(
"Query parsing error: {} (client: {})",
error, client_identifier
);
None
}
},
};
if let Some(ast) = ast {
let plugin_result = query_router.execute_plugins(&ast).await;
match plugin_result {
@@ -1237,8 +1279,6 @@ where
_ => (),
};
let _ = query_router.infer(&ast);
}
}
debug!("Sending query to server");
@@ -1290,7 +1330,7 @@ where
}
if query_router.query_parser_enabled() {
if let Ok(ast) = QueryRouter::parse(&message) {
if let Ok(ast) = query_router.parse(&message) {
if let Ok(output) = query_router.execute_plugins(&ast).await {
plugin_output = Some(output);
}