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..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::*; @@ -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" - ); -} 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()