From e85425dedae28988a8c4fe89e8a61dd9897b7193 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Fri, 7 Aug 2026 13:03:28 -0700 Subject: [PATCH 1/2] fix: schema-sharding function shouldnt be affected by omni writer check --- pgdog/src/frontend/router/parser/query/mod.rs | 8 +++-- .../router/parser/query/test/test_sharding.rs | 29 ++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 5844c8b0c..ad5d2125e 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -109,13 +109,17 @@ impl QueryParser { route.set_shard(context.shards_calculator.shard()); } + // Schema-based sharding. + let is_search_path = context.shards_calculator.is_search_path(); + // A write that only touches omnisharded tables must reach // every shard. A shard directive (a comment or SET) routing // it to one shard would silently diverge the table, so it's // an error. A pending lookup for a bare key means the same // directive is present with a cold cache, so it errors // without running the lookup. - if route.is_omnisharded() + if !is_search_path + && route.is_omnisharded() && route.is_write() && (matches!( route.shard_with_priority().source(), @@ -128,7 +132,7 @@ impl QueryParser { .pending_lookups .extend(std::mem::take(&mut context.bare_key_lookups)); - route.set_search_path_driven(context.shards_calculator.is_search_path()); + route.set_search_path_driven(is_search_path); route.set_pending_lookups(std::mem::take(&mut context.pending_lookups)); if let Some(role) = context.router_context.sticky.role { diff --git a/pgdog/src/frontend/router/parser/query/test/test_sharding.rs b/pgdog/src/frontend/router/parser/query/test/test_sharding.rs index c3fa0bad4..fe7d4e4c4 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_sharding.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_sharding.rs @@ -3,7 +3,8 @@ use std::ops::Deref; use crate::config::config; use crate::frontend::Command; -use crate::frontend::router::parser::{Cache, Shard}; +use crate::frontend::router::parser::{Cache, Shard, route::ShardSource}; +use crate::net::parameter::ParameterValue; use super::setup::{QueryParserTest, *}; @@ -426,6 +427,32 @@ fn test_comment_key_errors_on_omnisharded_write() { assert!(command.route().shard().is_direct()); } +/// A search_path route takes precedence over a comment key whose lookup +/// missed the cache, so an omnisharded write remains pinned to the schema's +/// shard instead of being rejected as a directive-driven write. +#[test] +fn test_search_path_allows_omnisharded_write_with_pending_comment_key() { + let tables = lookup_rule_tables(); + let mut test = QueryParserTest::new() + .with_sharded_tables(tables) + .with_param("search_path", ParameterValue::String("shard_0".into())); + + let command = test.execute(vec![ + Query::new( + "/* pgdog_sharding_key: 'org_child' */ INSERT INTO organizations (id, name) VALUES ('org_child', 'child')", + ) + .into(), + ]); + + assert!(command.route().is_omnisharded()); + assert!(command.route().is_search_path_driven()); + assert_eq!(command.route().shard(), &Shard::Direct(0)); + assert_eq!( + command.route().shard_with_priority().source(), + &ShardSource::SearchPath("shard_0".into()) + ); +} + /// SET pgdog.sharding_key errors on omnisharded writes the same way. #[test] fn test_set_key_errors_on_omnisharded_write() { From bd585d3d47a1728a217e3c411cfba68b9d4268e1 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Fri, 7 Aug 2026 13:10:46 -0700 Subject: [PATCH 2/2] fix: more omni checks --- .../client/query_engine/route_query.rs | 7 +--- pgdog/src/frontend/router/parser/query/mod.rs | 18 ++++----- .../router/parser/query/test/test_sharding.rs | 40 +++++++++++++++++-- pgdog/src/frontend/router/parser/route.rs | 18 +++++++++ 4 files changed, 63 insertions(+), 20 deletions(-) diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 1271cd37a..0bb668bcd 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -194,14 +194,9 @@ impl QueryEngine { // Make sure we don't send an omni write to a direct-to-shard route. // This will cause omni data inconsistency. fn is_omnishard_unsafe(backend: &Connection, command: &Command, cluster: &Cluster) -> bool { - command.route().is_omnisharded() - && command.route().is_write() + command.route().requires_full_shard_coverage() && backend.connected() // FIXME(lev): I wish there was a way to say >0 and bool { + self.is_omnisharded() && self.is_write() && !self.is_search_path_driven() + } + /// Return true if this route requires result set manipulation to /// return correct results. /// @@ -783,4 +791,14 @@ mod test { shards.push(ShardWithPriority::new_set(Shard::Direct(4))); assert_eq!(shards.shard().deref(), &Shard::Direct(3)); } + + #[test] + fn test_omnisharded_write_coverage_exempts_search_path_routes() { + let mut route = + Route::write(ShardWithPriority::new_table_omni(Shard::All)).with_omnisharded(true); + assert!(route.requires_full_shard_coverage()); + + route.set_search_path_driven(true); + assert!(!route.requires_full_shard_coverage()); + } }