Rewrite bug fixes - old cache key, more efficient update (#516)

* Rewrite bug fixes - old cache key, more efficient update

* Update pgdog/src/frontend/prepared_statements/mod.rs

Co-authored-by: Copilot <[email protected]>

* Update pgdog/src/frontend/client_request.rs

Co-authored-by: Lev Kokotov <[email protected]>

* Skip removing and readding cache item, just modify it in place

* skip cloning name if passed by reference

* skip cloning name if passed by reference

* format

* cache key fixes

* Update cache removal test to handle close_unused(0) clearing

---------

Co-authored-by: Copilot <[email protected]>
Co-authored-by: Lev Kokotov <[email protected]>

Author: 3 people

integration/rust/tests/integration/prepared.rs

@@ -79,6 +79,136 @@ async fn test_prepared_cache() {
     });
 }
 
+#[tokio::test]
+async fn test_prepared_cache_respects_limit() {
+    let admin = admin_sqlx().await;
+
+    // Start from a clean state so results aren't influenced by previous tests.
+    admin.execute("RECONNECT").await.unwrap();
+    admin
+        .execute("SET prepared_statements_limit TO 100")
+        .await
+        .unwrap();
+
+    // Run the average helper query a few times to populate the cache.
+    for _ in 0..3 {
+        let pools = connections_sqlx().await;
+        for pool in &pools {
+            sqlx::query("/* test_prepared_cache_rust */ SELECT $1")
+                .bind(5)
+                .fetch_one(pool)
+                .await
+                .unwrap();
+        }
+
+        for pool in pools {
+            pool.close().await;
+        }
+    }
+
+    let mut prepared = admin.fetch_all("SHOW PREPARED").await.unwrap();
+    prepared.retain(|row| {
+        row.get::<String, &str>("statement")
+            .contains("/* test_prepared_cache_rust")
+    });
+    assert_eq!(
+        prepared.len(),
+        2,
+        "expected the original statement plus its helper"
+    );
+
+    // Tighten the cache limit and ensure unused statements are evicted.
+    admin
+        .execute("SET prepared_statements_limit TO 1")
+        .await
+        .unwrap();
+
+    let mut prepared = admin.fetch_all("SHOW PREPARED").await.unwrap();
+    prepared.retain(|row| {
+        row.get::<String, &str>("statement")
+            .contains("/* test_prepared_cache_rust")
+    });
+    assert!(
+        prepared.len() <= 1,
+        "expected helper statements to be evicted when limit drops"
+    );
+
+    admin.execute("RELOAD").await.unwrap();
+}
+
+#[tokio::test]
+async fn test_prepared_cache_helper_evicted_on_close() {
+    let admin = admin_sqlx().await;
+    admin.execute("RECONNECT").await.unwrap();
+    admin
+        .execute("SET prepared_statements_limit TO 100")
+        .await
+        .unwrap();
+
+    let mut pools = connections_sqlx().await;
+    let sharded = pools.remove(1);
+    let primary = pools.remove(0);
+
+    // Build a deterministic dataset per shard to trigger AVG rewrite.
+    for shard in [0, 1] {
+        let drop = format!(
+            "/* pgdog_shard: {} */ DROP TABLE IF EXISTS avg_helper_cleanup",
+            shard
+        );
+        sharded.execute(drop.as_str()).await.ok();
+
+        let create = format!(
+            "/* pgdog_shard: {} */ CREATE TABLE avg_helper_cleanup(price DOUBLE PRECISION)",
+            shard
+        );
+        sharded.execute(create.as_str()).await.unwrap();
+    }
+
+    sharded
+        .execute("/* pgdog_shard: 0 */ INSERT INTO avg_helper_cleanup(price) VALUES (10.0), (14.0)")
+        .await
+        .unwrap();
+    sharded
+        .execute("/* pgdog_shard: 1 */ INSERT INTO avg_helper_cleanup(price) VALUES (18.0), (22.0)")
+        .await
+        .unwrap();
+
+    sqlx::query("/* test_avg_helper_cleanup */ SELECT AVG(price) FROM avg_helper_cleanup")
+        .fetch_one(&sharded)
+        .await
+        .unwrap();
+
+    // Clean up tables so subsequent tests start fresh.
+    for shard in [0, 1] {
+        let drop = format!(
+            "/* pgdog_shard: {} */ DROP TABLE IF EXISTS avg_helper_cleanup",
+            shard
+        );
+        sharded.execute(drop.as_str()).await.ok();
+    }
+
+    sharded.close().await;
+    primary.close().await;
+
+    let prepared = admin
+        .fetch_all("SHOW PREPARED")
+        .await
+        .unwrap()
+        .into_iter()
+        .filter(|row| {
+            row.get::<String, &str>("statement")
+                .contains("test_avg_helper_cleanup")
+        })
+        .collect::<Vec<_>>();
+
+    assert!(
+        prepared.is_empty(),
+        "helper rewrite statements should be evicted once the connection closes"
+    );
+
+    admin.execute("RELOAD").await.unwrap();
+}
+
 #[tokio::test]
 async fn test_prepard_cache_eviction() {
     let admin = admin_sqlx().await;

pgdog/src/frontend/client_request.rs

@@ -180,11 +180,7 @@ impl ClientRequest {
         for message in self.messages.iter_mut() {
             if let ProtocolMessage::Parse(parse) = message {
                 parse.set_query(query);
-                let name = parse.name().to_owned();
-                let _ = prepared.update_query(&name, query);
-                if !plan.is_noop() {
-                    prepared.set_rewrite_plan(&name, plan.clone());
-                }
+                prepared.update_and_set_rewrite_plan(&parse.name(), query, plan.clone());
                 updated = true;
             }
         }

pgdog/src/frontend/prepared_statements/global_cache.rs

@@ -17,8 +17,11 @@ fn global_name(counter: usize) -> String {
 pub struct Statement {
     parse: Parse,
     row_description: Option<RowDescription>,
+    #[allow(dead_code)]
     version: usize,
     rewrite_plan: Option<RewritePlan>,
+    cache_key: CacheKey,
+    evict_on_close: bool,
 }
 
 impl MemoryUsage for Statement {
@@ -30,8 +33,8 @@ impl MemoryUsage for Statement {
             } else {
                 0
             }
-            // Rewrite plans are small; treat as zero-cost for now.
-            + 0
+            + self.cache_key.memory_usage()
+            + self.evict_on_close.memory_usage()
     }
 }
 
@@ -41,11 +44,7 @@ impl Statement {
     }
 
     fn cache_key(&self) -> CacheKey {
-        CacheKey {
-            query: self.parse.query_ref(),
-            data_types: self.parse.data_types_ref(),
-            version: self.version,
-        }
+        self.cache_key.clone()
     }
 }
 
@@ -147,14 +146,14 @@ impl GlobalCache {
             let name = global_name(self.counter);
             let parse = parse.rename(&name);
 
-            let parse_key = CacheKey {
+            let cache_key = CacheKey {
                 query: parse.query_ref(),
                 data_types: parse.data_types_ref(),
                 version: 0,
             };
 
             self.statements.insert(
-                parse_key,
+                cache_key.clone(),
                 CachedStmt {
                     counter: self.counter,
                     used: 1,
@@ -168,6 +167,8 @@ impl GlobalCache {
                     row_description: None,
                     version: 0,
                     rewrite_plan: None,
+                    cache_key,
+                    evict_on_close: false,
                 },
             );
 
@@ -205,6 +206,8 @@ impl GlobalCache {
                 row_description: None,
                 version: self.versions,
                 rewrite_plan: None,
+                cache_key: key,
+                evict_on_close: false,
             },
         );
 
@@ -221,27 +224,27 @@ impl GlobalCache {
         }
     }
 
-    pub fn update_query(&mut self, name: &str, sql: &str) -> bool {
+    pub fn update_and_set_rewrite_plan(
+        &mut self,
+        name: &str,
+        sql: &str,
+        plan: RewritePlan,
+    ) -> bool {
         if let Some(statement) = self.names.get_mut(name) {
-            let old_key = statement.cache_key();
-            let cached = self.statements.remove(&old_key);
             statement.parse.set_query(sql);
-            let new_key = statement.cache_key();
-            if let Some(entry) = cached {
-                self.statements.insert(new_key, entry);
+            if !plan.is_noop() {
+                statement.evict_on_close = !plan.helpers().is_empty();
+                statement.rewrite_plan = Some(plan);
+            } else {
+                statement.evict_on_close = false;
+                statement.rewrite_plan = None;
             }
             true
         } else {
             false
         }
     }
 
-    pub fn set_rewrite_plan(&mut self, name: &str, plan: RewritePlan) {
-        if let Some(statement) = self.names.get_mut(name) {
-            statement.rewrite_plan = Some(plan);
-        }
-    }
-
     pub fn rewrite_plan(&self, name: &str) -> Option<RewritePlan> {
         self.names.get(name).and_then(|s| s.rewrite_plan.clone())
     }
@@ -290,27 +293,34 @@ impl GlobalCache {
 
     /// Close prepared statement.
     pub fn close(&mut self, name: &str, capacity: usize) -> bool {
-        let used = if let Some(stmt) = self.names.get(name) {
-            if let Some(stmt) = self.statements.get_mut(&stmt.cache_key()) {
-                stmt.used = stmt.used.saturating_sub(1);
-                stmt.used > 0
-            } else {
-                false
+        if let Some(statement) = self.names.get(name) {
+            let key = statement.cache_key();
+            let mut used_remaining = None;
+
+            if let Some(entry) = self.statements.get_mut(&key) {
+                entry.used = entry.used.saturating_sub(1);
+                used_remaining = Some(entry.used);
+                if entry.used == 0 && (statement.evict_on_close || self.len() > capacity) {
+                    self.remove(name);
+                    return true;
+                }
             }
-        } else {
-            false
-        };
 
-        if !used && self.len() > capacity {
-            self.remove(name);
-            true
-        } else {
-            false
+            return used_remaining.map(|u| u > 0).unwrap_or(false);
         }
+
+        false
     }
 
     /// Close all unused statements exceeding capacity.
     pub fn close_unused(&mut self, capacity: usize) -> usize {
+        if capacity == 0 {
+            let removed = self.statements.len();
+            self.statements.clear();
+            self.names.clear();
+            return removed;
+        }
+
         let mut remove = self.statements.len() as i64 - capacity as i64;
         let mut to_remove = vec![];
         for stmt in self.statements.values() {
@@ -409,7 +419,16 @@ mod test {
             names.push(name);
         }
 
-        assert_eq!(cache.close_unused(0), 0);
+        assert_eq!(cache.close_unused(0), 25);
+        assert!(cache.is_empty());
+
+        names.clear();
+        for stmt in 0..25 {
+            let parse = Parse::named("__sqlx_1", format!("SELECT {}", stmt));
+            let (new, name) = cache.insert(&parse);
+            assert!(new);
+            names.push(name);
+        }
 
         for name in &names[0..5] {
             assert!(!cache.close(name, 25)); // Won't close because
@@ -422,4 +441,48 @@ mod test {
         assert_eq!(cache.close_unused(19), 0);
         assert_eq!(cache.len(), 20);
     }
+
+    #[test]
+    fn test_close_unused_zero_clears_all_entries() {
+        let mut cache = GlobalCache::default();
+
+        for idx in 0..5 {
+            let parse = Parse::named("test", format!("SELECT {}", idx));
+            let (_is_new, _name) = cache.insert(&parse);
+        }
+
+        assert!(cache.len() > 0);
+
+        let removed = cache.close_unused(0);
+        assert_eq!(removed, 5);
+        assert!(cache.is_empty());
+    }
+
+    #[test]
+    fn test_update_query_reuses_cache_key() {
+        let mut cache = GlobalCache::default();
+        let parse = Parse::named("__sqlx_1", "SELECT 1");
+        let (is_new, name) = cache.insert(&parse);
+        assert!(is_new);
+
+        assert!(cache.update_and_set_rewrite_plan(
+            &name,
+            "SELECT 1 ORDER BY 1",
+            RewritePlan::default()
+        ));
+
+        let key = cache
+            .statements()
+            .keys()
+            .next()
+            .expect("statement key missing");
+        assert_eq!(key.query().unwrap(), "SELECT 1");
+        assert_eq!(cache.query(&name).unwrap(), "SELECT 1 ORDER BY 1");
+
+        let parse_again = Parse::named("__sqlx_2", "SELECT 1");
+        let (is_new_again, reused_name) = cache.insert(&parse_again);
+        assert!(!is_new_again);
+        assert_eq!(reused_name, name);
+        assert_eq!(cache.len(), 1);
+    }
 }