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

Fix broken custom config

Browse Source
This commit is contained in:
Lev Kokotov
2022-09-12 15:58:11 -04:00
parent efd6b2edae
commit bca5318d5e
1 changed files with 25 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/query_router.rs
View File

@@ -55,6 +55,8 @@ pub struct QueryRouter {
/// Include the primary into the replica pool for reads. /// Include the primary into the replica pool for reads.
primary_reads_enabled: bool, primary_reads_enabled: bool,
set_manually: bool,
/// Pool configuration. /// Pool configuration.
pool_settings: PoolSettings, pool_settings: PoolSettings,
} }
@@ -97,6 +99,7 @@ impl QueryRouter {
active_role: None, active_role: None,
query_parser_enabled: false, query_parser_enabled: false,
primary_reads_enabled: false, primary_reads_enabled: false,
set_manually: false,
pool_settings: PoolSettings::default(), pool_settings: PoolSettings::default(),
} }
} }
@@ -104,9 +107,12 @@ impl QueryRouter {
/// Pool settings can change because of a config reload. /// Pool settings can change because of a config reload.
pub fn update_pool_settings(&mut self, pool_settings: PoolSettings) { pub fn update_pool_settings(&mut self, pool_settings: PoolSettings) {
self.pool_settings = pool_settings; self.pool_settings = pool_settings;
if !self.set_manually {
self.query_parser_enabled = self.pool_settings.query_parser_enabled; self.query_parser_enabled = self.pool_settings.query_parser_enabled;
self.primary_reads_enabled = self.pool_settings.primary_reads_enabled; self.primary_reads_enabled = self.pool_settings.primary_reads_enabled;
} }
}
/// Try to parse a command and execute it. /// Try to parse a command and execute it.
pub fn try_execute_command(&mut self, mut buf: BytesMut) -> Option<(Command, String)> { pub fn try_execute_command(&mut self, mut buf: BytesMut) -> Option<(Command, String)> {
@@ -207,6 +213,8 @@ impl QueryRouter {
} }
Command::SetServerRole => { Command::SetServerRole => {
self.set_manually = true;
self.active_role = match value.to_ascii_lowercase().as_ref() { self.active_role = match value.to_ascii_lowercase().as_ref() {
"primary" => { "primary" => {
self.query_parser_enabled = false; self.query_parser_enabled = false;
@@ -230,7 +238,7 @@ impl QueryRouter {
"default" => { "default" => {
self.active_role = self.pool_settings.default_role; self.active_role = self.pool_settings.default_role;
self.query_parser_enabled = self.query_parser_enabled; self.query_parser_enabled = self.pool_settings.query_parser_enabled;
self.active_role self.active_role
} }
@@ -239,6 +247,8 @@ impl QueryRouter {
} }
Command::SetPrimaryReads => { Command::SetPrimaryReads => {
self.set_manually = true;
if value == "on" { if value == "on" {
debug!("Setting primary reads to on"); debug!("Setting primary reads to on");
self.primary_reads_enabled = true; self.primary_reads_enabled = true;
@@ -400,7 +410,7 @@ mod test {
} }
#[test] #[test]
fn test__replica() { fn test_infer_role_and_shard_replica() {
QueryRouter::setup(); QueryRouter::setup();
let mut qr = QueryRouter::new(); let mut qr = QueryRouter::new();
assert!(qr.try_execute_command(simple_query("SET SERVER ROLE TO 'auto'")) != None); assert!(qr.try_execute_command(simple_query("SET SERVER ROLE TO 'auto'")) != None);
@@ -409,17 +419,20 @@ mod test {
assert!(qr.try_execute_command(simple_query("SET PRIMARY READS TO off")) != None); assert!(qr.try_execute_command(simple_query("SET PRIMARY READS TO off")) != None);
let queries = vec![ let queries = vec![
simple_query("SELECT * FROM items WHERE id = 5"), simple_query("SELECT * FROM items WHERE id = 4"),
simple_query( simple_query(
"SELECT id, name, value FROM items INNER JOIN prices ON item.id = prices.item_id", "SELECT id, name, value FROM items INNER JOIN prices ON item.id = prices.item_id",
), ),
simple_query("WITH t AS (SELECT * FROM items) SELECT * FROM t"), simple_query("WITH t AS (SELECT * FROM items) SELECT * FROM t"),
]; ];
for query in queries { let shards = vec![0, 0, 0];
for (idx, query) in queries.iter().enumerate() {
// It's a recognized query // It's a recognized query
assert!(qr.infer_role_and_shard(query)); assert!(qr.infer_role_and_shard(query.clone()));
assert_eq!(qr.role(), Some(Role::Replica)); assert_eq!(qr.role(), Some(Role::Replica));
assert_eq!(qr.shard(), shards[idx]);
} }
} }
@@ -643,7 +656,7 @@ mod test {
assert!(qr.query_parser_enabled()); assert!(qr.query_parser_enabled());
let query = simple_query("SET SERVER ROLE TO 'default'"); let query = simple_query("SET SERVER ROLE TO 'default'");
assert!(qr.try_execute_command(query) != None); assert!(qr.try_execute_command(query) != None);
assert!(qr.query_parser_enabled()); assert!(!qr.query_parser_enabled());
} }
#[test] #[test]
@@ -656,7 +669,8 @@ mod test {
user: crate::config::User::default(), user: crate::config::User::default(),
default_role: Some(Role::Replica), default_role: Some(Role::Replica),
query_parser_enabled: true, query_parser_enabled: true,
primary_reads_enabled: false, primary_reads_enabled: true,
sharding_key_regex: None,
sharding_function: ShardingFunction::PgBigintHash, sharding_function: ShardingFunction::PgBigintHash,
}; };
let mut qr = QueryRouter::new(); let mut qr = QueryRouter::new();
@@ -670,8 +684,8 @@ mod test {
assert_eq!(qr.active_role, None); assert_eq!(qr.active_role, None);
assert_eq!(qr.active_shard, None); assert_eq!(qr.active_shard, None);
assert_eq!(qr.query_parser_enabled, false); assert_eq!(qr.query_parser_enabled, true);
assert_eq!(qr.primary_reads_enabled, false); assert_eq!(qr.primary_reads_enabled, true);
let q1 = simple_query("SET SERVER ROLE TO 'primary'"); let q1 = simple_query("SET SERVER ROLE TO 'primary'");
assert!(qr.try_execute_command(q1) != None); assert!(qr.try_execute_command(q1) != None);