pgdog.git / summary / log / commit / refs
commit 856e3293ef024bfc10e7d014ecff6ba1eb338ecc
Author: Lev Kokotov <levkk@users.noreply.github.com>
Commit: GitHub <noreply@github.com>
Date: Tue Dec 02 18:48:57 2025 +0000
fix: schema sharding (#641)
- fix: `SELECT` queries were routed to catch-all schema when a more
specific one is available
- fix: `search_path` schema wasn't taking priority in DDL statements
- fix: remove duplicate `search_path` schema check
pgdog/src/frontend/router/parser/query/ddl.rs | 6 +-
pgdog/src/frontend/router/parser/query/delete.rs | 2 +-
pgdog/src/frontend/router/parser/query/mod.rs | 41 +------
.../router/parser/query/schema_sharding.rs | 38 +++----
pgdog/src/frontend/router/parser/query/select.rs | 32 +++---
pgdog/src/frontend/router/parser/query/shared.rs | 34 +++---
pgdog/src/frontend/router/parser/query/update.rs | 2 +-
pgdog/src/frontend/router/sharding/mod.rs | 2 +
pgdog/src/frontend/router/sharding/schema.rs | 123 +++++++++++++++++++++
pgdog/src/unique_id.rs | 2 +-
10 files changed, 177 insertions(+), 105 deletions(-)
diff --git a/pgdog/src/frontend/router/parser/query/ddl.rs b/pgdog/src/frontend/router/parser/query/ddl.rs
index d45db691..910d7cab 100644
--- a/pgdog/src/frontend/router/parser/query/ddl.rs
+++ b/pgdog/src/frontend/router/parser/query/ddl.rs
@@ -12,10 +12,8 @@ impl QueryParser {
let mut command = Self::shard_ddl(node, &context.sharding_schema)?;
if let Command::Query(ref mut route) = command {
- if route.shard().all() {
- if let Some(shard) = self.check_search_path_for_shard(context)? {
- route.set_shard_mut(shard);
- }
+ if let Some(shard) = self.check_search_path_for_shard(context)? {
+ route.set_shard_mut(shard);
}
}
diff --git a/pgdog/src/frontend/router/parser/query/delete.rs b/pgdog/src/frontend/router/parser/query/delete.rs
index f2be1e7e..a414ddf8 100644
--- a/pgdog/src/frontend/router/parser/query/delete.rs
+++ b/pgdog/src/frontend/router/parser/query/delete.rs
@@ -40,7 +40,7 @@ impl QueryParser {
context.router_context.bind,
&mut self.explain_recorder,
)?;
- let shard = Self::converge(shards, ConvergeAlgorithm::default());
+ 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/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs
index b94a8f22..fd91b4d2 100644
--- a/pgdog/src/frontend/router/parser/query/mod.rs
+++ b/pgdog/src/frontend/router/parser/query/mod.rs
@@ -68,6 +68,7 @@ pub struct QueryParser {
shard: Shard,
// Plugin read override.
plugin_output: PluginOutput,
+ // Record explain output.
explain_recorder: Option<ExplainRecorder>,
}
@@ -382,46 +383,6 @@ impl QueryParser {
// with the fingerprint.
//
if route.shard().all() {
- // Check search_path for schema.
- let search_path = context.router_context.params.get("search_path");
-
- // Quick inline function to shard query
- // based on schema in search_path.
- fn shard_from_search_path(
- search_path: &str,
- context: &QueryParserContext<'_>,
- query_parser: &mut QueryParser,
- route: &mut Route,
- ) {
- let schema = Schema::from(search_path);
-
- if let Some(schema) = context.sharding_schema.schemas.get(Some(schema)) {
- let shard: Shard = schema.shard().into();
-
- if let Some(recorder) = query_parser.recorder_mut() {
- recorder.record_entry(
- Some(shard.clone()),
- format!("matched schema {} in search_path", schema.name()),
- );
- }
- route.set_shard_mut(shard);
- }
- }
-
- match search_path {
- Some(ParameterValue::String(search_path)) => {
- shard_from_search_path(search_path, context, self, route);
- }
-
- Some(ParameterValue::Tuple(search_paths)) => {
- for schema in search_paths {
- shard_from_search_path(schema, context, self, route);
- }
- }
-
- None => (),
- }
-
let databases = databases();
// Only fingerprint the query if some manual queries are configured.
// Otherwise, we're wasting time parsing SQL.
diff --git a/pgdog/src/frontend/router/parser/query/schema_sharding.rs b/pgdog/src/frontend/router/parser/query/schema_sharding.rs
index 9131e196..1a564259 100644
--- a/pgdog/src/frontend/router/parser/query/schema_sharding.rs
+++ b/pgdog/src/frontend/router/parser/query/schema_sharding.rs
@@ -1,3 +1,5 @@
+use crate::frontend::router::sharding::SchemaSharder;
+
use super::*;
impl QueryParser {
@@ -12,16 +14,17 @@ impl QueryParser {
// Check search_path for schema.
let search_path = context.router_context.params.get("search_path");
+ let mut schema_sharder = SchemaSharder::default();
match search_path {
Some(ParameterValue::String(search_path)) => {
let schema = Schema::from(search_path.as_str());
- if let Some(schema) = context.sharding_schema.schemas.get(Some(schema)) {
- let shard: Shard = schema.shard().into();
+ schema_sharder.resolve(Some(schema), &context.sharding_schema.schemas);
+ if let Some((shard, schema)) = schema_sharder.get() {
if let Some(recorder) = self.recorder_mut() {
recorder.record_entry(
Some(shard.clone()),
- format!("matched schema \"{}\" in search_path", schema.name()),
+ format!("matched schema {} in search_path", schema),
);
}
return Ok(Some(shard));
@@ -29,30 +32,19 @@ impl QueryParser {
}
Some(ParameterValue::Tuple(search_paths)) => {
- let mut candidates = vec![];
-
- for (idx, schema) in search_paths.iter().enumerate() {
+ for schema in search_paths {
let schema = Schema::from(schema.as_str());
- if let Some(schema) = context.sharding_schema.schemas.get(Some(schema)) {
- let shard: Shard = schema.shard().into();
- let catch_all = schema.is_default();
- candidates.push((shard, catch_all, idx));
- }
+ schema_sharder.resolve(Some(schema), &context.sharding_schema.schemas);
}
- // false < true
- // Catch-all schemas go first, more qualified ones go last.
- candidates.sort_by_key(|cand| !cand.1);
- if let Some(candidate) = candidates.pop() {
- if let Some(schema) = search_paths.get(candidate.2) {
- if let Some(recorder) = self.recorder_mut() {
- recorder.record_entry(
- Some(candidate.0.clone()),
- format!("matched schema mult \"{}\" in search_path", schema),
- );
- }
+ if let Some((shard, schema)) = schema_sharder.get() {
+ if let Some(recorder) = self.recorder_mut() {
+ recorder.record_entry(
+ Some(shard.clone()),
+ format!("matched schema {} in search_path", schema),
+ );
}
- return Ok(Some(candidate.0));
+ return Ok(Some(shard));
}
}
diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs
index 22c24c0c..3c1524ac 100644
--- a/pgdog/src/frontend/router/parser/query/select.rs
+++ b/pgdog/src/frontend/router/parser/query/select.rs
@@ -1,5 +1,6 @@
-use crate::frontend::router::parser::{
- cache::CachedAst, from_clause::FromClause, where_clause::TablesSource,
+use crate::frontend::router::{
+ parser::{cache::CachedAst, from_clause::FromClause, where_clause::TablesSource},
+ sharding::SchemaSharder,
};
use super::*;
@@ -65,23 +66,18 @@ impl QueryParser {
}
// Schema-based sharding.
+ let mut schema_sharder = SchemaSharder::default();
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.clone()),
- format!("SELECT matched schema {}", schema.name()),
- );
- }
- }
+ let schema = table.schema();
+ schema_sharder.resolve(schema, &context.sharding_schema.schemas);
+ }
+ if let Some((shard, schema)) = schema_sharder.get() {
+ if let Some(recorder) = self.recorder_mut() {
+ recorder.record_entry(
+ Some(shard.clone()),
+ format!("SELECT matched schema {}", schema),
+ );
}
-
- // Converge to the first direct shard.
- let shard = Self::converge(shards.clone(), ConvergeAlgorithm::FirstDirect);
- shards = HashSet::new();
shards.insert(shard);
}
@@ -109,7 +105,7 @@ impl QueryParser {
}
}
- let shard = Self::converge(shards, ConvergeAlgorithm::default());
+ 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()?;
diff --git a/pgdog/src/frontend/router/parser/query/shared.rs b/pgdog/src/frontend/router/parser/query/shared.rs
index cd79a600..803f9bd4 100644
--- a/pgdog/src/frontend/router/parser/query/shared.rs
+++ b/pgdog/src/frontend/router/parser/query/shared.rs
@@ -14,13 +14,13 @@ pub(super) enum ConvergeAlgorithm {
impl QueryParser {
/// Converge to a single route given multiple shards.
- pub(super) fn converge(shards: HashSet<Shard>, algorithm: ConvergeAlgorithm) -> Shard {
+ pub(super) fn converge(shards: &HashSet<Shard>, algorithm: ConvergeAlgorithm) -> Shard {
let shard = if shards.len() == 1 {
shards.iter().next().cloned().unwrap()
} else {
let mut multi = vec![];
let mut all = false;
- for shard in &shards {
+ for shard in shards {
match shard {
Shard::All => {
all = true;
@@ -189,10 +189,10 @@ mod tests {
fn single_direct_returns_itself() {
let shards = HashSet::from([Shard::Direct(5)]);
- let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
assert_eq!(result, Shard::Direct(5));
- let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::Direct(5));
}
@@ -200,10 +200,10 @@ mod tests {
fn single_all_returns_itself() {
let shards = HashSet::from([Shard::All]);
- let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
assert_eq!(result, Shard::All);
- let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::All);
}
@@ -211,10 +211,10 @@ mod tests {
fn single_multi_returns_itself() {
let shards = HashSet::from([Shard::Multi(vec![1, 2, 3])]);
- let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
assert_eq!(result, Shard::Multi(vec![1, 2, 3]));
- let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::Multi(vec![1, 2, 3]));
}
@@ -222,7 +222,7 @@ mod tests {
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);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
match result {
Shard::Multi(mut v) => {
v.sort();
@@ -236,7 +236,7 @@ mod tests {
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);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert!(
matches!(result, Shard::Direct(1) | Shard::Direct(2)),
"expected Direct(1) or Direct(2), got {:?}",
@@ -248,7 +248,7 @@ mod tests {
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);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
assert_eq!(result, Shard::All);
}
@@ -256,7 +256,7 @@ mod tests {
fn all_present_first_direct_returns_direct() {
let shards = HashSet::from([Shard::All, Shard::Direct(1)]);
- let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::Direct(1));
}
@@ -264,10 +264,10 @@ mod tests {
fn empty_set_returns_all() {
let shards = HashSet::new();
- let result = QueryParser::converge(shards.clone(), ConvergeAlgorithm::AllFirstElseMulti);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
assert_eq!(result, Shard::All);
- let result = QueryParser::converge(shards, ConvergeAlgorithm::FirstDirect);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::All);
}
@@ -275,7 +275,7 @@ mod tests {
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);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::AllFirstElseMulti);
match result {
Shard::Multi(mut v) => {
v.sort();
@@ -289,7 +289,7 @@ mod tests {
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);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::Direct(3));
}
@@ -297,7 +297,7 @@ mod tests {
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);
+ let result = QueryParser::converge(&shards, ConvergeAlgorithm::FirstDirect);
assert_eq!(result, Shard::All);
}
}
diff --git a/pgdog/src/frontend/router/parser/query/update.rs b/pgdog/src/frontend/router/parser/query/update.rs
index 0db7321f..351890e0 100644
--- a/pgdog/src/frontend/router/parser/query/update.rs
+++ b/pgdog/src/frontend/router/parser/query/update.rs
@@ -62,7 +62,7 @@ impl QueryParser {
context.router_context.bind,
&mut self.explain_recorder,
)?;
- let shard = Self::converge(shards, ConvergeAlgorithm::default());
+ 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/sharding/mod.rs b/pgdog/src/frontend/router/sharding/mod.rs
index a5ecc49e..deac0e61 100644
--- a/pgdog/src/frontend/router/sharding/mod.rs
+++ b/pgdog/src/frontend/router/sharding/mod.rs
@@ -22,6 +22,7 @@ pub mod list;
pub mod mapping;
pub mod operator;
pub mod range;
+pub mod schema;
pub mod tables;
#[cfg(test)]
pub mod test;
@@ -33,6 +34,7 @@ pub use context_builder::*;
pub use error::Error;
pub use hasher::Hasher;
pub use operator::*;
+pub use schema::SchemaSharder;
pub use tables::*;
pub use value::*;
pub use vector::{Centroids, Distance};
diff --git a/pgdog/src/frontend/router/sharding/schema.rs b/pgdog/src/frontend/router/sharding/schema.rs
new file mode 100644
index 00000000..893d1e1e
--- /dev/null
+++ b/pgdog/src/frontend/router/sharding/schema.rs
@@ -0,0 +1,123 @@
+use crate::{
+ backend::replication::ShardedSchemas,
+ frontend::router::parser::{Schema, Shard},
+};
+
+#[derive(Debug, Default, Clone)]
+pub struct SchemaSharder {
+ catch_all: bool,
+ current: Option<Shard>,
+ schema: Option<String>,
+}
+
+impl SchemaSharder {
+ /// Resolve current schema.
+ pub fn resolve(&mut self, schema: Option<Schema<'_>>, schemas: &ShardedSchemas) {
+ let check = schemas.get(schema);
+ if let Some(schema) = check {
+ let catch_all = schema.is_default();
+ let set =
+ catch_all && self.current.is_none() || self.catch_all || self.current.is_none();
+ if set {
+ self.current = Some(schema.shard().into());
+ self.catch_all = catch_all;
+ self.schema = Some(schema.name().to_owned());
+ }
+ }
+ }
+
+ pub fn get(&self) -> Option<(Shard, &str)> {
+ if let Some(current) = self.current.as_ref() {
+ if let Some(schema) = self.schema.as_ref() {
+ return Some((current.clone(), schema.as_str()));
+ }
+ }
+
+ None
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use pgdog_config::sharding::ShardedSchema;
+
+ #[test]
+ fn test_catch_all_chosen_only_when_no_specific_match() {
+ // Create a catch-all schema (no name) that routes to shard 0
+ let catch_all = ShardedSchema {
+ database: "test".to_string(),
+ name: None, // This makes it a catch-all/default
+ shard: 0,
+ all: false,
+ };
+
+ // Create a specific schema "sales" that routes to shard 1
+ let sales_schema = ShardedSchema {
+ database: "test".to_string(),
+ name: Some("sales".to_string()),
+ shard: 1,
+ all: false,
+ };
+
+ let schemas = ShardedSchemas::new(vec![catch_all, sales_schema]);
+
+ // Test 1: When we resolve "sales", we should get shard 1 (specific match, not catch-all)
+ let mut sharder = SchemaSharder::default();
+ sharder.resolve(Some(Schema { name: "sales" }), &schemas);
+ let result = sharder.get();
+ assert_eq!(
+ result.as_ref().map(|(s, _)| s.clone()),
+ Some(Shard::Direct(1))
+ );
+ assert_eq!(result.as_ref().map(|(_, name)| *name), Some("sales"));
+ assert!(!sharder.catch_all);
+
+ // Test 2: When we resolve "unknown", we should get shard 0 (catch-all)
+ let mut sharder = SchemaSharder::default();
+ sharder.resolve(Some(Schema { name: "unknown" }), &schemas);
+ let result = sharder.get();
+ assert_eq!(
+ result.as_ref().map(|(s, _)| s.clone()),
+ Some(Shard::Direct(0))
+ );
+ assert_eq!(result.as_ref().map(|(_, name)| *name), Some("*"));
+ assert!(sharder.catch_all);
+
+ // Test 3: When we resolve None (no schema specified), we should get shard 0 (catch-all)
+ let mut sharder = SchemaSharder::default();
+ sharder.resolve(None, &schemas);
+ let result = sharder.get();
+ assert_eq!(
+ result.as_ref().map(|(s, _)| s.clone()),
+ Some(Shard::Direct(0))
+ );
+ assert_eq!(result.as_ref().map(|(_, name)| *name), Some("*"));
+ assert!(sharder.catch_all);
+
+ // Test 4: When catch-all is resolved first, then specific match,
+ // the specific match should override the catch-all
+ let mut sharder = SchemaSharder::default();
+ sharder.resolve(Some(Schema { name: "unknown" }), &schemas); // catch-all first
+ sharder.resolve(Some(Schema { name: "sales" }), &schemas); // specific second
+ let result = sharder.get();
+ assert_eq!(
+ result.as_ref().map(|(s, _)| s.clone()),
+ Some(Shard::Direct(1))
+ );
+ assert_eq!(result.as_ref().map(|(_, name)| *name), Some("sales"));
+ assert!(!sharder.catch_all);
+
+ // Test 5: When specific match is resolved first, catch-all should NOT override
+ let mut sharder = SchemaSharder::default();
+ sharder.resolve(Some(Schema { name: "sales" }), &schemas); // specific first
+ sharder.resolve(Some(Schema { name: "unknown" }), &schemas); // catch-all second
+ let result = sharder.get();
+ assert_eq!(
+ result.as_ref().map(|(s, _)| s.clone()),
+ Some(Shard::Direct(1))
+ );
+ assert_eq!(result.as_ref().map(|(_, name)| *name), Some("sales"));
+ assert!(!sharder.catch_all);
+ }
+}
diff --git a/pgdog/src/unique_id.rs b/pgdog/src/unique_id.rs
index 65d77a04..660e2814 100644
--- a/pgdog/src/unique_id.rs
+++ b/pgdog/src/unique_id.rs
@@ -249,7 +249,7 @@ mod test {
let extracted_node2 = (id2 >> NODE_SHIFT) & MAX_NODE_ID;
assert_eq!(extracted_node2, node);
- assert!(extracted_seq2 >= 1); // Sequence incremented (or time advanced and reset to 0)
+ assert!(matches!(extracted_seq2, 1 | 0)); // Sequence incremented (or time advanced and reset to 0)
}
#[tokio::test]
[parent: 0ac6914038d0]