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

More query router commands; settings last until changed again; docs (#25)

Browse Source 
* readme

* touch up docs

* stuff

* refactor query router

* remove unused

* less verbose

* docs

* no link

* method rename
This commit is contained in:
Lev Kokotov
2022-02-19 08:57:24 -08:00
committed by GitHub
parent bbacb9cf01
commit a556ec1c43
5 changed files with 579 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
src/client.rs
View File

@@ -14,7 +14,7 @@ use crate::constants::*;
use crate::errors::Error;
use crate::messages::*;
use crate::pool::{ClientServerMap, ConnectionPool};
use crate::query_router::QueryRouter;
use crate::query_router::{Command, QueryRouter};
use crate::server::Server;
use crate::stats::Reporter;
@@ -198,29 +198,50 @@ impl Client {
// SET SHARDING KEY TO 'bigint';
let mut message = read_message(&mut self.read).await?;
// Parse for special select shard command.
// SET SHARDING KEY TO 'bigint';
if query_router.select_shard(message.clone()) {
custom_protocol_response_ok(
// Handle all custom protocol commands here.
match query_router.try_execute_command(message.clone()) {
// Normal query
None => {
if query_router.query_parser_enabled() && query_router.role() == None {
query_router.infer_role(message.clone());
}
}
Some((Command::SetShard, _)) | Some((Command::SetShardingKey, _)) => {
custom_protocol_response_ok(&mut self.write, &format!("SET SHARD")).await?;
continue;
}
Some((Command::SetServerRole, _)) => {
custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
continue;
}
Some((Command::ShowServerRole, value)) => {
show_response(&mut self.write, "server role", &value).await?;
continue;
}
Some((Command::ShowShard, value)) => {
show_response(&mut self.write, "shard", &value).await?;
continue;
}
};
// Make sure we selected a valid shard.
if query_router.shard() >= pool.shards() {
error_response(
&mut self.write,
&format!("SET SHARD TO {}", query_router.shard()),
&format!(
"shard '{}' is more than configured '{}'",
query_router.shard(),
pool.shards()
),
)
.await?;
continue;
}
// Parse for special server role selection command.
// SET SERVER ROLE TO '(primary|replica)';
if query_router.select_role(message.clone()) {
custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
continue;
}
// Attempt to parse the query to determine where it should go
if query_router.query_parser_enabled() && query_router.role() == None {
query_router.infer_role(message.clone());
}
// Grab a server from the pool: the client issued a regular query.
let connection = match pool.get(query_router.shard(), query_router.role()).await {
Ok(conn) => conn,
@@ -228,7 +249,6 @@ impl Client {
println!(">> Could not get connection from pool: {:?}", err);
error_response(&mut self.write, "could not get connection from the pool")
.await?;
query_router.reset();
continue;
}
};
@@ -310,9 +330,6 @@ impl Client {
if self.transaction_mode {
// Report this client as idle.
self.stats.client_idle();
query_router.reset();
break;
}
}
@@ -395,9 +412,6 @@ impl Client {
if self.transaction_mode {
self.stats.client_idle();
query_router.reset();
break;
}
}
@@ -431,8 +445,7 @@ impl Client {
self.stats.transaction();
if self.transaction_mode {
query_router.reset();
self.stats.client_idle();
break;
}
}