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()); + } }