fix: schema-sharding function shouldnt be affected by omni writer check

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 <n in one shot.
             && backend.connected_servers() < cluster.shards().len()
-            // Schema-based sharding intentionally routes an omnisharded table to the
-            // single shard selected by search_path; only accidental partial routing is unsafe.
-            // FIXME(lev): Encode the required execution scope in Route instead of inferring it here.
-            && !command.route().is_search_path_driven()
     }
 
     // Caller switched shards mid-transaction and the transaction is pinned

pgdog/src/frontend/router/parser/query/mod.rs

@@ -109,14 +109,17 @@ impl QueryParser {
                     route.set_shard(context.shards_calculator.shard());
                 }
 
-                // 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()
-                    && route.is_write()
+                // Schema-based sharding.
+                let is_search_path = context.shards_calculator.is_search_path();
+                route.set_search_path_driven(is_search_path);
+
+                // A non-schema-routed 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 is deferred when
+                // search_path currently controls the route; after the lookup is
+                // resolved, the second routing pass performs this check again.
+                if route.requires_full_shard_coverage()
                     && (matches!(
                         route.shard_with_priority().source(),
                         ShardSource::Comment | ShardSource::Set
@@ -128,7 +131,6 @@ 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_pending_lookups(std::mem::take(&mut context.pending_lookups));
 
                 if let Some(role) = context.router_context.sticky.role {

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,64 @@ fn test_comment_key_errors_on_omnisharded_write() {
     assert!(command.route().shard().is_direct());
 }
 
+/// A search_path route defers rejecting a comment key whose lookup missed the
+/// cache. The pending lookup is resolved before execution, and the second
+/// routing pass can then decide whether the comment overrides search_path.
+#[test]
+fn test_search_path_defers_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().pending_lookups().len(), 1);
+    assert_eq!(command.route().shard(), &Shard::Direct(0));
+    assert_eq!(
+        command.route().shard_with_priority().source(),
+        &ShardSource::SearchPath("shard_0".into())
+    );
+}
+
+/// Once the deferred comment lookup resolves, routing checks the omnisharded
+/// write again. The higher-priority comment now selects one shard, so the write
+/// must be rejected instead of partially updating the table.
+#[test]
+fn test_resolved_comment_key_errors_on_search_path_omnisharded_write() {
+    use crate::frontend::router::parser::Error;
+    use crate::frontend::router::sharding::{LookupTable, ResolvedLookups};
+
+    let key = LookupTable {
+        schema: None,
+        name: None,
+        column: "organization_id".into(),
+    };
+    let mut resolved = ResolvedLookups::default();
+    resolved.insert(key, "org_child".into(), "org_root".into());
+
+    let mut test = QueryParserTest::new()
+        .with_sharded_tables(lookup_rule_tables())
+        .with_param("search_path", ParameterValue::String("shard_0".into()))
+        .with_resolved_lookups(resolved);
+
+    let result = test.try_execute(vec![
+        Query::new(
+            "/* pgdog_sharding_key: 'org_child' */ INSERT INTO organizations (id, name) VALUES ('org_child', 'child')",
+        )
+        .into(),
+    ]);
+
+    assert!(matches!(result, Err(Error::OmniWriteWithDirective)));
+}
+
 /// SET pgdog.sharding_key errors on omnisharded writes the same way.
 #[test]
 fn test_set_key_errors_on_omnisharded_write() {

pgdog/src/frontend/router/parser/route.rs

@@ -283,6 +283,14 @@ impl Route {
         self.search_path_driven
     }
 
+    /// Whether an omnisharded write must reach every shard to remain consistent.
+    ///
+    /// Schema-based sharding intentionally limits the write to the shard selected
+    /// by `search_path`; every other omnisharded write requires full coverage.
+    pub(crate) fn requires_full_shard_coverage(&self) -> 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());
+    }
 }