pgdog.git / summary / log / commit / refs
commit 6762a4e2d128b2b26e0d7b89618bd755fa1737f4
Author: Lev Kokotov <levkk@users.noreply.github.com>
Commit: GitHub <noreply@github.com>
Date: Fri Nov 21 19:57:32 2025 +0000
Omni routing and change to convergence for schema sharding (#620)
* Omni routing and change to convergence for schema sharding
* save
* Fix test
* Add config test
* Rename
integration/d_plus/pgdog.toml | 1 +
pgdog-config/src/core.rs | 65 +++++++++-
pgdog-config/src/sharding.rs | 8 ++
pgdog/src/backend/pool/cluster.rs | 13 +-
.../replication/logical/subscriber/context.rs | 1 +
pgdog/src/backend/replication/sharded_tables.rs | 18 ++-
pgdog/src/frontend/client/mod.rs | 5 +-
pgdog/src/frontend/client/query_engine/context.rs | 6 +
.../frontend/client/query_engine/route_query.rs | 1 +
pgdog/src/frontend/router/cli.rs | 1 +
pgdog/src/frontend/router/context.rs | 4 +
pgdog/src/frontend/router/parser/cache.rs | 2 +-
pgdog/src/frontend/router/parser/query/delete.rs | 3 +-
pgdog/src/frontend/router/parser/query/explain.rs | 4 +-
pgdog/src/frontend/router/parser/query/select.rs | 33 +++--
pgdog/src/frontend/router/parser/query/shared.rs | 143 ++++++++++++++++++++-
pgdog/src/frontend/router/parser/query/show.rs | 4 +-
pgdog/src/frontend/router/parser/query/test.rs | 83 +++++++++---
pgdog/src/frontend/router/parser/query/update.rs | 3 +-
pgdog/src/frontend/router/parser/route.rs | 4 +
20 files changed, 361 insertions(+), 41 deletions(-)
diff --git a/integration/d_plus/pgdog.toml b/integration/d_plus/pgdog.toml
index 1f51cd65..a7f5ca71 100644
--- a/integration/d_plus/pgdog.toml
+++ b/integration/d_plus/pgdog.toml
@@ -23,3 +23,4 @@ tables = [
"pg_namespace",
"pg_am"
]
+sticky = true
diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs
index faf39fc3..a43c3443 100644
--- a/pgdog-config/src/core.rs
+++ b/pgdog-config/src/core.rs
@@ -5,7 +5,9 @@ use std::path::PathBuf;
use tracing::{info, warn};
use crate::sharding::ShardedSchema;
-use crate::{EnumeratedDatabase, Memory, PassthoughAuth, PreparedStatements, RewriteMode};
+use crate::{
+ EnumeratedDatabase, Memory, OmnishardedTable, PassthoughAuth, PreparedStatements, RewriteMode,
+};
use super::database::Database;
use super::error::Error;
@@ -227,7 +229,7 @@ impl Config {
tables
}
- pub fn omnisharded_tables(&self) -> HashMap<String, Vec<String>> {
+ pub fn omnisharded_tables(&self) -> HashMap<String, Vec<OmnishardedTable>> {
let mut tables = HashMap::new();
for table in &self.omnisharded_tables {
@@ -235,7 +237,10 @@ impl Config {
.entry(table.database.clone())
.or_insert_with(Vec::new);
for t in &table.tables {
- entry.push(t.clone());
+ entry.push(OmnishardedTable {
+ name: t.clone(),
+ sticky_routing: table.sticky,
+ });
}
}
@@ -647,4 +652,58 @@ password = "users_admin_password"
);
assert!(config_and_users.users.admin.is_none());
}
+
+ #[test]
+ fn test_omnisharded_tables() {
+ let source = r#"
+[general]
+host = "0.0.0.0"
+port = 6432
+
+[[databases]]
+name = "db1"
+host = "127.0.0.1"
+port = 5432
+
+[[databases]]
+name = "db2"
+host = "127.0.0.1"
+port = 5433
+
+[[omnisharded_tables]]
+database = "db1"
+tables = ["table_a", "table_b"]
+
+[[omnisharded_tables]]
+database = "db1"
+tables = ["table_c"]
+sticky = true
+
+[[omnisharded_tables]]
+database = "db2"
+tables = ["table_x"]
+"#;
+
+ let config: Config = toml::from_str(source).unwrap();
+
+ assert_eq!(config.omnisharded_tables.len(), 3);
+
+ let tables = config.omnisharded_tables();
+
+ assert_eq!(tables.len(), 2);
+
+ let db1_tables = tables.get("db1").unwrap();
+ assert_eq!(db1_tables.len(), 3);
+ assert_eq!(db1_tables[0].name, "table_a");
+ assert!(!db1_tables[0].sticky_routing);
+ assert_eq!(db1_tables[1].name, "table_b");
+ assert!(!db1_tables[1].sticky_routing);
+ assert_eq!(db1_tables[2].name, "table_c");
+ assert!(db1_tables[2].sticky_routing);
+
+ let db2_tables = tables.get("db2").unwrap();
+ assert_eq!(db2_tables.len(), 1);
+ assert_eq!(db2_tables[0].name, "table_x");
+ assert!(!db2_tables[0].sticky_routing);
+ }
}
diff --git a/pgdog-config/src/sharding.rs b/pgdog-config/src/sharding.rs
index 51a0b214..35dcf3bb 100644
--- a/pgdog-config/src/sharding.rs
+++ b/pgdog-config/src/sharding.rs
@@ -177,6 +177,14 @@ impl From<String> for FlexibleType {
pub struct OmnishardedTables {
pub database: String,
pub tables: Vec<String>,
+ #[serde(default)]
+ pub sticky: bool,
+}
+
+#[derive(PartialEq, Debug, Clone, Default)]
+pub struct OmnishardedTable {
+ pub name: String,
+ pub sticky_routing: bool,
}
/// Queries with manual routing rules.
diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs
index be6cb76a..db273dca 100644
--- a/pgdog/src/backend/pool/cluster.rs
+++ b/pgdog/src/backend/pool/cluster.rs
@@ -538,6 +538,8 @@ impl Cluster {
mod test {
use std::{sync::Arc, time::Duration};
+ use pgdog_config::OmnishardedTable;
+
use crate::{
backend::{
pool::{Address, Config, PoolConfig, ShardConfig},
@@ -597,7 +599,16 @@ mod test {
hasher: Hasher::Postgres,
..Default::default()
}],
- vec!["sharded_omni".into()],
+ vec![
+ OmnishardedTable {
+ name: "sharded_omni".into(),
+ sticky_routing: false,
+ },
+ OmnishardedTable {
+ name: "sharded_omni_sticky".into(),
+ sticky_routing: true,
+ },
+ ],
),
shards,
identifier,
diff --git a/pgdog/src/backend/replication/logical/subscriber/context.rs b/pgdog/src/backend/replication/logical/subscriber/context.rs
index 32f984a6..fe54c756 100644
--- a/pgdog/src/backend/replication/logical/subscriber/context.rs
+++ b/pgdog/src/backend/replication/logical/subscriber/context.rs
@@ -56,6 +56,7 @@ impl<'a> StreamContext<'a> {
&mut self.prepared_statements,
&self.params,
None,
+ 1,
)?)
}
}
diff --git a/pgdog/src/backend/replication/sharded_tables.rs b/pgdog/src/backend/replication/sharded_tables.rs
index 805f50c5..3cf54541 100644
--- a/pgdog/src/backend/replication/sharded_tables.rs
+++ b/pgdog/src/backend/replication/sharded_tables.rs
@@ -1,15 +1,20 @@
//! Tables sharded in the database.
+use pgdog_config::OmnishardedTable;
+
use crate::{
config::{DataType, ShardedTable},
frontend::router::sharding::Mapping,
net::messages::Vector,
};
-use std::{collections::HashSet, sync::Arc};
+use std::{
+ collections::{HashMap, HashSet},
+ sync::Arc,
+};
#[derive(Default, Debug)]
struct Inner {
tables: Vec<ShardedTable>,
- omnisharded: HashSet<String>,
+ omnisharded: HashMap<String, bool>, // Name <-> sticky routing
/// This is set only if we have the same sharding scheme
/// across all tables, i.e., 3 tables with the same data type
/// and list/range/hash function.
@@ -46,7 +51,7 @@ impl From<&[ShardedTable]> for ShardedTables {
}
impl ShardedTables {
- pub fn new(tables: Vec<ShardedTable>, omnisharded_tables: Vec<String>) -> Self {
+ pub fn new(tables: Vec<ShardedTable>, omnisharded_tables: Vec<OmnishardedTable>) -> Self {
let mut common_mapping = HashSet::new();
for table in &tables {
common_mapping.insert((
@@ -69,7 +74,10 @@ impl ShardedTables {
Self {
inner: Arc::new(Inner {
tables,
- omnisharded: omnisharded_tables.into_iter().collect(),
+ omnisharded: omnisharded_tables
+ .into_iter()
+ .map(|table| (table.name, table.sticky_routing))
+ .collect(),
common_mapping,
}),
}
@@ -79,7 +87,7 @@ impl ShardedTables {
&self.inner.tables
}
- pub fn omnishards(&self) -> &HashSet<String> {
+ pub fn omnishards(&self) -> &HashMap<String, bool> {
&self.inner.omnisharded
}
diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs
index a3650b86..7ecb8019 100644
--- a/pgdog/src/frontend/client/mod.rs
+++ b/pgdog/src/frontend/client/mod.rs
@@ -4,6 +4,7 @@ use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
+use rand::{thread_rng, Rng};
use timeouts::Timeouts;
use tokio::{select, spawn, time::timeout};
use tracing::{debug, enabled, error, info, trace, Level as LogLevel};
@@ -28,7 +29,6 @@ use crate::state::State;
use crate::stats::memory::MemoryUsage;
use crate::util::user_database_from_params;
-// pub mod counter;
pub mod query_engine;
pub mod timeouts;
@@ -50,6 +50,7 @@ pub struct Client {
client_request: ClientRequest,
stream_buffer: MessageBuffer,
passthrough_password: Option<String>,
+ omni_sticky_index: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -315,6 +316,7 @@ impl Client {
stream_buffer: MessageBuffer::new(config.config.memory.message_buffer),
shutdown: false,
passthrough_password,
+ omni_sticky_index: thread_rng().gen_range(1..usize::MAX),
}))
}
@@ -342,6 +344,7 @@ impl Client {
stream_buffer: MessageBuffer::new(4096),
shutdown: false,
passthrough_password: None,
+ omni_sticky_index: 1,
}
}
diff --git a/pgdog/src/frontend/client/query_engine/context.rs b/pgdog/src/frontend/client/query_engine/context.rs
index dbff3d61..87ede236 100644
--- a/pgdog/src/frontend/client/query_engine/context.rs
+++ b/pgdog/src/frontend/client/query_engine/context.rs
@@ -1,3 +1,5 @@
+use rand::{thread_rng, Rng};
+
use crate::{
backend::pool::{connection::mirror::Mirror, stats::MemoryStats},
frontend::{
@@ -34,6 +36,8 @@ pub struct QueryEngineContext<'a> {
pub(super) admin: bool,
/// Executing rollback statement.
pub(super) rollback: bool,
+ /// Omnisharded modulo.
+ pub(super) omni_sticky_index: usize,
}
impl<'a> QueryEngineContext<'a> {
@@ -53,6 +57,7 @@ impl<'a> QueryEngineContext<'a> {
admin: client.admin,
requests_left: 0,
rollback: false,
+ omni_sticky_index: client.omni_sticky_index,
}
}
@@ -77,6 +82,7 @@ impl<'a> QueryEngineContext<'a> {
admin: false,
requests_left: 0,
rollback: false,
+ omni_sticky_index: thread_rng().gen_range(1..usize::MAX),
}
}
diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs
index 4b229e50..6dd9382f 100644
--- a/pgdog/src/frontend/client/query_engine/route_query.rs
+++ b/pgdog/src/frontend/client/query_engine/route_query.rs
@@ -50,6 +50,7 @@ impl QueryEngine {
context.prepared_statements,
context.params,
context.transaction,
+ context.omni_sticky_index,
)?;
match self.router.query(router_context) {
Ok(cmd) => {
diff --git a/pgdog/src/frontend/router/cli.rs b/pgdog/src/frontend/router/cli.rs
index 9f445d08..711eb351 100644
--- a/pgdog/src/frontend/router/cli.rs
+++ b/pgdog/src/frontend/router/cli.rs
@@ -55,6 +55,7 @@ impl RouterCli {
&mut stmt,
&mut params,
None,
+ 1,
)?)?;
result.push(cmd);
}
diff --git a/pgdog/src/frontend/router/context.rs b/pgdog/src/frontend/router/context.rs
index fa58fdb3..f430cc0f 100644
--- a/pgdog/src/frontend/router/context.rs
+++ b/pgdog/src/frontend/router/context.rs
@@ -25,6 +25,8 @@ pub struct RouterContext<'a> {
pub executable: bool,
/// Two-pc enabled
pub two_pc: bool,
+ /// Sticky omnisharded index.
+ pub omni_sticky_index: usize,
}
impl<'a> RouterContext<'a> {
@@ -34,6 +36,7 @@ impl<'a> RouterContext<'a> {
stmt: &'a mut PreparedStatements,
params: &'a Parameters,
transaction: Option<TransactionType>,
+ omni_sticky_index: usize,
) -> Result<Self, Error> {
let query = buffer.query()?;
let bind = buffer.parameters()?;
@@ -49,6 +52,7 @@ impl<'a> RouterContext<'a> {
copy_mode,
executable: buffer.executable(),
two_pc: cluster.two_pc_enabled(),
+ omni_sticky_index,
})
}
diff --git a/pgdog/src/frontend/router/parser/cache.rs b/pgdog/src/frontend/router/parser/cache.rs
index 27bccc8e..3d0d5b16 100644
--- a/pgdog/src/frontend/router/parser/cache.rs
+++ b/pgdog/src/frontend/router/parser/cache.rs
@@ -68,7 +68,7 @@ impl Deref for CachedAst {
impl CachedAst {
/// Create new cache entry from pg_query's AST.
- fn new(query: &str, schema: &ShardingSchema) -> std::result::Result<Self, super::Error> {
+ pub fn new(query: &str, schema: &ShardingSchema) -> std::result::Result<Self, super::Error> {
let ast = parse(query).map_err(super::Error::PgQuery)?;
let (shard, role) = comment(query, schema)?;
diff --git a/pgdog/src/frontend/router/parser/query/delete.rs b/pgdog/src/frontend/router/parser/query/delete.rs
index 60138269..ef98b426 100644
--- a/pgdog/src/frontend/router/parser/query/delete.rs
+++ b/pgdog/src/frontend/router/parser/query/delete.rs
@@ -1,5 +1,6 @@
use crate::frontend::router::parser::where_clause::TablesSource;
+use super::shared::ConvergeAlgorithm;
use super::*;
impl QueryParser {
@@ -35,7 +36,7 @@ impl QueryParser {
context.router_context.bind,
&mut self.explain_recorder,
)?;
- let shard = Self::converge(shards);
+ let shard = Self::converge(shards, ConvergeAlgorithm::default());
if let Some(recorder) = self.recorder_mut() {
recorder.record_entry(
Some(shard.clone()),
diff --git a/pgdog/src/frontend/router/parser/query/explain.rs b/pgdog/src/frontend/router/parser/query/explain.rs
index 878b6c62..8b38261a 100644
--- a/pgdog/src/frontend/router/parser/query/explain.rs
+++ b/pgdog/src/frontend/router/parser/query/explain.rs
@@ -75,7 +75,7 @@ mod tests {
let mut stmts = PreparedStatements::default();
let params = Parameters::default();
- let ctx = RouterContext::new(&buffer, &cluster, &mut stmts, ¶ms, None).unwrap();
+ let ctx = RouterContext::new(&buffer, &cluster, &mut stmts, ¶ms, None, 1).unwrap();
match QueryParser::default().parse(ctx).unwrap().clone() {
Command::Query(route) => route,
@@ -102,7 +102,7 @@ mod tests {
let mut stmts = PreparedStatements::default();
let params = Parameters::default();
- let ctx = RouterContext::new(&buffer, &cluster, &mut stmts, ¶ms, None).unwrap();
+ let ctx = RouterContext::new(&buffer, &cluster, &mut stmts, ¶ms, None, 1).unwrap();
match QueryParser::default().parse(ctx).unwrap().clone() {
Command::Query(route) => route,
diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs
index ac7e97c2..2b1ca19d 100644
--- a/pgdog/src/frontend/router/parser/query/select.rs
+++ b/pgdog/src/frontend/router/parser/query/select.rs
@@ -3,6 +3,7 @@ use crate::frontend::router::parser::{
};
use super::*;
+use shared::ConvergeAlgorithm;
impl QueryParser {
/// Handle SELECT statement.
@@ -63,15 +64,21 @@ impl QueryParser {
for table in cached_ast.tables() {
if let Some(schema) = context.sharding_schema.schemas.get(table.schema()) {
let shard: Shard = schema.shard().into();
+
if shards.insert(shard.clone()) {
if let Some(recorder) = self.recorder_mut() {
recorder.record_entry(
- Some(shard),
+ Some(shard.clone()),
format!("SELECT matched schema {}", schema.name()),
);
}
}
}
+
+ // Converge to the first direct shard.
+ let shard = Self::converge(shards.clone(), ConvergeAlgorithm::FirstDirect);
+ shards = HashSet::new();
+ shards.insert(shard);
}
// Shard by vector in ORDER BY clause.
@@ -98,7 +105,7 @@ impl QueryParser {
}
}
- let shard = Self::converge(shards);
+ let shard = Self::converge(shards, ConvergeAlgorithm::default());
let aggregates = Aggregate::parse(stmt)?;
let limit = LimitClause::new(stmt, context.router_context.bind).limit_offset()?;
let distinct = Distinct::new(stmt).distinct()?;
@@ -108,16 +115,26 @@ impl QueryParser {
// Omnisharded tables check.
if query.is_all_shards() {
let tables = from_clause.tables();
+ let mut sticky = false;
let omni = tables.iter().all(|table| {
- context
- .sharding_schema
- .tables
- .omnishards()
- .contains(table.name)
+ let is_sticky = context.sharding_schema.tables.omnishards().get(table.name);
+
+ if let Some(is_sticky) = is_sticky {
+ if *is_sticky {
+ sticky = true;
+ }
+ true
+ } else {
+ false
+ }
});
if omni {
- let shard = round_robin::next() % context.shards;
+ let shard = if sticky {
+ context.router_context.omni_sticky_index
+ } else {
+ round_robin::next()
+ } % context.shards;
query.set_shard_mut(shard);
diff --git a/pgdog/src/frontend/router/parser/query/shared.rs b/pgdog/src/frontend/router/parser/query/shared.rs
index 26e3fe6f..f5c61c38 100644
--- a/pgdog/src/frontend/router/parser/query/shared.rs
+++ b/pgdog/src/frontend/router/parser/query/shared.rs
@@ -1,9 +1,20 @@
use super::{explain_trace::ExplainRecorder, *};
use std::string::String as StdString;
+#[derive(Debug, Clone, Default, Copy, PartialEq)]
+pub(super) enum ConvergeAlgorithm {
+ // Take the first direct shard we find
+ FirstDirect,
+ // If All is present, make it cross-shard.
+ // If multiple shards are present, make it multi.
+ // Else, make it direct.
+ #[default]
+ AllFirstElseMulti,
+}
+
impl QueryParser {
/// Converge to a single route given multiple shards.
- pub(super) fn converge(shards: HashSet<Shard>) -> Shard {
+ pub(super) fn converge(shards: HashSet<Shard>, algorithm: ConvergeAlgorithm) -> Shard {
let shard = if shards.len() == 1 {
shards.iter().next().cloned().unwrap()
} else {
@@ -19,6 +30,14 @@ impl QueryParser {
Shard::Multi(m) => multi.extend(m),
};
}
+
+ if algorithm == ConvergeAlgorithm::FirstDirect {
+ let direct = shards.iter().find(|shard| shard.is_direct());
+ if let Some(direct) = direct {
+ return direct.clone();
+ }
+ }
+
if all || shards.is_empty() {
Shard::All
} else {
@@ -142,3 +161,125 @@ fn record_column<F>(
recorder.record_entry(shard, description);
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::collections::HashSet;
+
+ #[test]
+ fn single_direct_returns_itself() {
+ let shards = HashSet::from([Shard::Direct(5)]);
+
+ let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ assert_eq!(result, Shard::Direct(5));
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::Direct(5));
+ }
+
+ #[test]
+ fn single_all_returns_itself() {
+ let shards = HashSet::from([Shard::All]);
+
+ let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ assert_eq!(result, Shard::All);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::All);
+ }
+
+ #[test]
+ fn single_multi_returns_itself() {
+ let shards = HashSet::from([Shard::Multi(vec![1, 2, 3])]);
+
+ let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ assert_eq!(result, Shard::Multi(vec![1, 2, 3]));
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::Multi(vec![1, 2, 3]));
+ }
+
+ #[test]
+ fn multiple_direct_all_first_else_multi_returns_multi() {
+ let shards = HashSet::from([Shard::Direct(1), Shard::Direct(2)]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::AllFirstElseMulti);
+ match result {
+ Shard::Multi(mut v) => {
+ v.sort();
+ assert_eq!(v, vec![1, 2]);
+ }
+ other => panic!("expected Multi, got {:?}", other),
+ }
+ }
+
+ #[test]
+ fn multiple_direct_first_direct_returns_one_direct() {
+ let shards = HashSet::from([Shard::Direct(1), Shard::Direct(2)]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert!(
+ matches!(result, Shard::Direct(1) | Shard::Direct(2)),
+ "expected Direct(1) or Direct(2), got {:?}",
+ result
+ );
+ }
+
+ #[test]
+ fn all_present_all_first_else_multi_returns_all() {
+ let shards = HashSet::from([Shard::All, Shard::Direct(1)]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::AllFirstElseMulti);
+ assert_eq!(result, Shard::All);
+ }
+
+ #[test]
+ fn all_present_first_direct_returns_direct() {
+ let shards = HashSet::from([Shard::All, Shard::Direct(1)]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::Direct(1));
+ }
+
+ #[test]
+ fn empty_set_returns_all() {
+ let shards = HashSet::new();
+
+ let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ assert_eq!(result, Shard::All);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::All);
+ }
+
+ #[test]
+ fn multi_and_direct_merge_into_multi() {
+ let shards = HashSet::from([Shard::Multi(vec![1, 2]), Shard::Direct(3)]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::AllFirstElseMulti);
+ match result {
+ Shard::Multi(mut v) => {
+ v.sort();
+ assert_eq!(v, vec![1, 2, 3]);
+ }
+ other => panic!("expected Multi, got {:?}", other),
+ }
+ }
+
+ #[test]
+ fn multi_and_direct_first_direct_returns_direct() {
+ let shards = HashSet::from([Shard::Multi(vec![1, 2]), Shard::Direct(3)]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::Direct(3));
+ }
+
+ #[test]
+ fn all_with_multi_first_direct_no_direct_returns_all() {
+ let shards = HashSet::from([Shard::All, Shard::Multi(vec![1, 2])]);
+
+ let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ assert_eq!(result, Shard::All);
+ }
+}
diff --git a/pgdog/src/frontend/router/parser/query/show.rs b/pgdog/src/frontend/router/parser/query/show.rs
index 1e34bc90..b146cfcf 100644
--- a/pgdog/src/frontend/router/parser/query/show.rs
+++ b/pgdog/src/frontend/router/parser/query/show.rs
@@ -38,7 +38,7 @@ mod test_show {
// First call
let query = "SHOW TRANSACTION ISOLATION LEVEL";
let buffer = ClientRequest::from(vec![Query::new(query).into()]);
- let context = RouterContext::new(&buffer, &c, &mut ps, &p, None).unwrap();
+ let context = RouterContext::new(&buffer, &c, &mut ps, &p, None, 1).unwrap();
let first = parser.parse(context).unwrap().clone();
let first_shard = first.route().shard();
@@ -47,7 +47,7 @@ mod test_show {
// Second call
let query = "SHOW TRANSACTION ISOLATION LEVEL";
let buffer = ClientRequest::from(vec![Query::new(query).into()]);
- let context = RouterContext::new(&buffer, &c, &mut ps, &p, None).unwrap();
+ let context = RouterContext::new(&buffer, &c, &mut ps, &p, None, 1).unwrap();
let second = parser.parse(context).unwrap().clone();
let second_shard = second.route().shard();
diff --git a/pgdog/src/frontend/router/parser/query/test.rs b/pgdog/src/frontend/router/parser/query/test.rs
index e197495c..88f5b6f5 100644
--- a/pgdog/src/frontend/router/parser/query/test.rs
+++ b/pgdog/src/frontend/router/parser/query/test.rs
@@ -70,6 +70,19 @@ fn lock_config_mode() -> MutexGuard<'static, ()> {
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
+fn parse_query(query: &str) -> Command {
+ let mut query_parser = QueryParser::default();
+ let client_request = ClientRequest::from(vec![Query::new(query).into()]);
+ let cluster = Cluster::new_test();
+ let mut stmt = PreparedStatements::default();
+ let params = Parameters::default();
+
+ let context =
+ RouterContext::new(&client_request, &cluster, &mut stmt, ¶ms, None, 1).unwrap();
+ let command = query_parser.parse(context).unwrap().clone();
+ command
+}
+
macro_rules! command {
($query:expr) => {{
command!($query, false)
@@ -87,8 +100,15 @@ macro_rules! command {
} else {
None
};
- let context =
- RouterContext::new(&client_request, &cluster, &mut stmt, ¶ms, transaction).unwrap();
+ let context = RouterContext::new(
+ &client_request,
+ &cluster,
+ &mut stmt,
+ ¶ms,
+ transaction,
+ 1,
+ )
+ .unwrap();
let command = query_parser.parse(context).unwrap().clone();
(command, query_parser)
@@ -131,6 +151,7 @@ macro_rules! query_parser {
&mut prep_stmts,
¶ms,
maybe_transaction,
+ 1,
)
.unwrap();
@@ -165,6 +186,7 @@ macro_rules! parse {
&mut PreparedStatements::default(),
&Parameters::default(),
None,
+ 1,
)
.unwrap(),
)
@@ -188,7 +210,7 @@ fn parse_with_parameters(query: &str, params: Parameters) -> Result<Command, Err
let client_request: ClientRequest = vec![Query::new(query).into()].into();
let cluster = Cluster::new_test();
let router_context =
- RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None).unwrap();
+ RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None, 1).unwrap();
QueryParser::default().parse(router_context)
}
@@ -209,6 +231,7 @@ fn parse_with_bind(query: &str, params: &[&[u8]]) -> Result<Command, Error> {
&mut prep_stmts,
¶meters,
None,
+ 1,
)
.unwrap();
@@ -324,11 +347,41 @@ fn test_prepared_stddev_rewrite_plan() {
#[test]
fn test_omni() {
- let q = "SELECT sharded_omni.* FROM sharded_omni WHERE sharded_omni.id = $1";
- let route = query!(q);
- assert!(matches!(route.shard(), Shard::Direct(_)));
- let (_, qp) = command!(q);
- assert!(!qp.in_transaction);
+ let mut omni_round_robin = HashSet::new();
+ let q = "SELECT sharded_omni.* FROM sharded_omni WHERE sharded_omni.id = 1";
+
+ for _ in 0..10 {
+ let command = parse_query(q);
+ match command {
+ Command::Query(query) => {
+ assert!(matches!(query.shard(), Shard::Direct(_)));
+ omni_round_robin.insert(query.shard().clone());
+ }
+
+ _ => {}
+ }
+ }
+
+ assert_eq!(omni_round_robin.len(), 2);
+
+ // Test sticky routing
+ let mut omni_sticky = HashSet::new();
+ let q =
+ "SELECT sharded_omni_sticky.* FROM sharded_omni_sticky WHERE sharded_omni_sticky.id = $1";
+
+ for _ in 0..10 {
+ let command = parse_query(q);
+ match command {
+ Command::Query(query) => {
+ assert!(matches!(query.shard(), Shard::Direct(_)));
+ omni_sticky.insert(query.shard().clone());
+ }
+
+ _ => {}
+ }
+ }
+
+ assert_eq!(omni_sticky.len(), 1);
// Test that sharded tables take priority.
let q = "
@@ -438,7 +491,7 @@ fn test_set() {
let params = Parameters::default();
let transaction = Some(TransactionType::ReadWrite);
let router_context =
- RouterContext::new(&buffer, &cluster, &mut prep_stmts, ¶ms, transaction).unwrap();
+ RouterContext::new(&buffer, &cluster, &mut prep_stmts, ¶ms, transaction, 1).unwrap();
let mut context = QueryParserContext::new(router_context);
for read_only in [true, false] {
@@ -519,7 +572,7 @@ fn update_sharding_key_errors_by_default() {
let client_request: ClientRequest = vec![Query::new(query).into()].into();
let cluster = Cluster::new_test();
let router_context =
- RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None).unwrap();
+ RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None, 1).unwrap();
let result = QueryParser::default().parse(router_context);
assert!(
@@ -539,7 +592,7 @@ fn update_sharding_key_ignore_mode_allows() {
let client_request: ClientRequest = vec![Query::new(query).into()].into();
let cluster = Cluster::new_test();
let router_context =
- RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None).unwrap();
+ RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None, 1).unwrap();
let command = QueryParser::default().parse(router_context).unwrap();
assert!(matches!(command, Command::Query(_)));
@@ -556,7 +609,7 @@ fn update_sharding_key_rewrite_mode_not_supported() {
let client_request: ClientRequest = vec![Query::new(query).into()].into();
let cluster = Cluster::new_test();
let router_context =
- RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None).unwrap();
+ RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None, 1).unwrap();
let result = QueryParser::default().parse(router_context);
assert!(
@@ -576,7 +629,7 @@ fn update_sharding_key_rewrite_plan_detected() {
let client_request: ClientRequest = vec![Query::new(query).into()].into();
let cluster = Cluster::new_test();
let router_context =
- RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None).unwrap();
+ RouterContext::new(&client_request, &cluster, &mut prep_stmts, ¶ms, None, 1).unwrap();
let command = QueryParser::default().parse(router_context).unwrap();
match command {
@@ -752,7 +805,7 @@ WHERE t2.account = (
.into();
let transaction = Some(TransactionType::ReadWrite);
let router_context =
- RouterContext::new(&buffer, &cluster, &mut prep_stmts, ¶ms, transaction).unwrap();
+ RouterContext::new(&buffer, &cluster, &mut prep_stmts, ¶ms, transaction, 1).unwrap();
let mut context = QueryParserContext::new(router_context);
let route = qp.query(&mut context).unwrap();
match route {
@@ -813,7 +866,7 @@ fn test_close_direct_one_shard() {
let params = Parameters::default();
let transaction = None;
- let context = RouterContext::new(&buf, &cluster, &mut pp, ¶ms, transaction).unwrap();
+ let context = RouterContext::new(&buf, &cluster, &mut pp, ¶ms, transaction, 1).unwrap();
let cmd = qp.parse(context).unwrap();
diff --git a/pgdog/src/frontend/router/parser/query/update.rs b/pgdog/src/frontend/router/parser/query/update.rs
index 8df7a0dd..55ca2ffe 100644
--- a/pgdog/src/frontend/router/parser/query/update.rs
+++ b/pgdog/src/frontend/router/parser/query/update.rs
@@ -9,6 +9,7 @@ use crate::{
};
use pg_query::protobuf::ColumnRef;
+use super::shared::ConvergeAlgorithm;
use super::*;
impl QueryParser {
@@ -57,7 +58,7 @@ impl QueryParser {
context.router_context.bind,
&mut self.explain_recorder,
)?;
- let shard = Self::converge(shards);
+ let shard = Self::converge(shards, ConvergeAlgorithm::default());
if let Some(recorder) = self.recorder_mut() {
recorder.record_entry(
Some(shard.clone()),
diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs
index 8b969893..af1799d1 100644
--- a/pgdog/src/frontend/router/parser/route.rs
+++ b/pgdog/src/frontend/router/parser/route.rs
@@ -35,6 +35,10 @@ impl Shard {
pub fn direct(shard: usize) -> Self {
Self::Direct(shard)
}
+
+ pub fn is_direct(&self) -> bool {
+ matches!(self, Self::Direct(_))
+ }
}
impl From<Option<usize>> for Shard {
[parent: 919e5ee59dd0]