From 0b96b463c3678f0935183e99c5984e1f0ac83419 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Wed, 7 Oct 2026 15:40:22 -0600 Subject: [PATCH 1/2] Cache queries regardless of routing comments Routing comments do not affect the rewriter in any way, and routing happens after rewriting occurs. There's no reason we need to avoid caching these --- .../router/parser/cache/cache_impl.rs | 15 +---- .../router/parser/query/test/test_comments.rs | 59 ------------------- 2 files changed, 3 insertions(+), 71 deletions(-) diff --git a/pgdog/src/frontend/router/parser/cache/cache_impl.rs b/pgdog/src/frontend/router/parser/cache/cache_impl.rs index 2219fe157..c057c7bb9 100644 --- a/pgdog/src/frontend/router/parser/cache/cache_impl.rs +++ b/pgdog/src/frontend/router/parser/cache/cache_impl.rs @@ -135,18 +135,9 @@ impl Cache { let parse_time = ast.stats.lock().parse_time; let mut guard = self.inner.lock(); - // Don't cache when a shard comment routed the query AND a rewrite - // was applied: the cache key is the comment-stripped body, so a - // subsequent uncommented lookup would hit this entry and receive an - // already-rewritten plan that was built against the commented - // (direct-shard) variant. - let cacheable = - query_and_comment.comment.shard.is_none() || ast.rewrite_plan.is_empty(); - if cacheable { - guard - .queries - .put(ast.query_without_comment.clone(), Arc::clone(&ast)); - } + guard + .queries + .put(ast.query_without_comment.clone(), Arc::clone(&ast)); guard.stats.misses += 1; guard.stats.parse_time += parse_time; Ok(ast) diff --git a/pgdog/src/frontend/router/parser/query/test/test_comments.rs b/pgdog/src/frontend/router/parser/query/test/test_comments.rs index d7b231029..2588e2208 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_comments.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_comments.rs @@ -41,62 +41,3 @@ fn test_comment_pgdog_shard_extended() { assert_eq!(command.route().shard(), &Shard::Direct(1234)); } - -/// A shard-commented query that requires a rewrite must not be stored in -/// the cache under the comment-stripped key. If it were, a subsequent -/// uncommented query with the same body would cache-hit and inherit the -/// rewrite plan built against the direct-shard variant. -/// -/// `pgdog.unique_id()` is a reliable rewrite trigger: the rewriter replaces -/// the call with a `$N::bigint` parameter in extended protocol. -#[test] -fn test_shard_comment_with_rewrite_not_cached() { - Cache::reset(); - let mut test = QueryParserTest::new(); - - let stripped = "SELECT pgdog.unique_id()"; - let commented = "/* pgdog_shard: 0 */ SELECT pgdog.unique_id()"; - - test.execute(vec![ - Parse::named("__shard_comment_rewrite", commented).into(), - Bind::new_statement("__shard_comment_rewrite").into(), - Execute::new().into(), - Sync.into(), - ]); - - let cached = Cache::queries(); - let poisoned = cached.keys().any(|k| **k == *stripped); - assert!( - !poisoned, - "shard-commented query with a non-empty rewrite plan must not be \ - cached under the stripped key; a subsequent uncommented lookup \ - would otherwise receive the direct-shard plan" - ); -} - -/// Complement to the previous test: a shard-commented query whose rewrite -/// plan is empty can safely share a cache entry with its uncommented -/// counterpart, since there's no plan to poison. Verifies the cacheable -/// check isn't blanket-dropping every shard-commented query. -#[test] -fn test_shard_comment_without_rewrite_is_cached() { - Cache::reset(); - let mut test = QueryParserTest::new(); - - let stripped = "SELECT 1"; - let commented = "/* pgdog_shard: 0 */ SELECT 1"; - - test.execute(vec![ - Parse::named("__shard_comment_no_rewrite", commented).into(), - Bind::new_statement("__shard_comment_no_rewrite").into(), - Execute::new().into(), - Sync.into(), - ]); - - let cached = Cache::queries(); - assert!( - cached.keys().any(|k| **k == *stripped), - "shard-commented query with an empty rewrite plan must still be \ - cached under the stripped key" - ); -} From b33a7a526f39c4bfb88798454a6ee0cb0220f521 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Wed, 7 Oct 2026 16:18:15 -0600 Subject: [PATCH 2/2] Fix dead code --- pgdog/src/frontend/router/parser/query/test/test_comments.rs | 2 +- pgdog/src/frontend/router/parser/rewrite/statement/plan.rs | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pgdog/src/frontend/router/parser/query/test/test_comments.rs b/pgdog/src/frontend/router/parser/query/test/test_comments.rs index 2588e2208..0c560eb8b 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_comments.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_comments.rs @@ -1,4 +1,4 @@ -use crate::frontend::router::parser::{Cache, Shard}; +use crate::frontend::router::parser::Shard; use super::setup::*; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 606a9977f..a176d06cc 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -40,6 +40,7 @@ impl BindParams { } } + #[cfg(test)] pub(super) fn is_original(&self) -> bool { matches!(self, Self::Original { .. }) } @@ -202,6 +203,7 @@ impl RewritePlan { /// True if the plan would not modify the query or its messages. /// `params` is purely informational (count of original `$N` placeholders) /// and doesn't count as a rewrite. + #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { self.bind_params.is_original() && self.stmt.is_none()