Cache queries regardless of routing comments

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)

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"
-    );
-}

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()