Dont rewrite anon prepared statements by default (#556)

* Dont rewrite anon prepared statements by default

* Fix test

Author: levkk

integration/go/go_pq/go_pq_test.go

@@ -81,6 +81,13 @@ func TestAuthenticationWithPassthrough(t *testing.T) {
 }
 
 func TestPqCrud(t *testing.T) {
+	adminConn, err := sql.Open("postgres", "postgres://admin:pgdog@127.0.0.1:6432/admin?sslmode=disable")
+	assert.Nil(t, err)
+	defer adminConn.Close()
+
+	_, err = adminConn.Exec("SET prepared_statements TO 'extended_anonymous'")
+	assert.Nil(t, err)
+
 	conns := PqConnections()
 
 	for _, conn := range conns {

integration/pgdog.toml

@@ -11,6 +11,7 @@ read_write_strategy = "aggressive"
 openmetrics_port = 9090
 openmetrics_namespace = "pgdog_"
 prepared_statements_limit = 500
+prepared_statements = "extended"
 expanded_explain = true
 # dns_ttl = 15_000
 query_cache_limit = 500

pgdog/src/admin/set.rs

@@ -91,6 +91,10 @@ impl Command for Set {
                     .close_unused(config.config.general.prepared_statements_limit);
             }
 
+            "prepared_statements" => {
+                config.config.general.prepared_statements = Self::from_json(&self.value)?;
+            }
+
             "cross_shard_disabled" => {
                 config.config.general.cross_shard_disabled = Self::from_json(&self.value)?;
             }

pgdog/src/config/core.rs

@@ -4,6 +4,8 @@ use std::fs::read_to_string;
 use std::path::PathBuf;
 use tracing::{info, warn};
 
+use crate::config::PreparedStatements;
+
 use super::database::Database;
 use super::error::Error;
 use super::general::General;
@@ -82,12 +84,12 @@ impl ConfigAndUsers {
     }
 
     /// Prepared statements are enabled.
-    pub fn prepared_statements(&self) -> bool {
+    pub fn prepared_statements(&self) -> PreparedStatements {
         // Disable prepared statements automatically in session mode
         if self.config.general.pooler_mode == PoolerMode::Session {
-            false
+            PreparedStatements::Disabled
         } else {
-            self.config.general.prepared_statements.enabled()
+            self.config.general.prepared_statements
         }
     }
 
@@ -378,31 +380,34 @@ column = "tenant_id"
         // Test transaction mode (default) - prepared statements should be enabled
         config.config.general.pooler_mode = PoolerMode::Transaction;
         config.config.general.prepared_statements = PreparedStatements::Extended;
-        assert!(
+        assert_eq!(
             config.prepared_statements(),
+            PreparedStatements::Extended,
             "Prepared statements should be enabled in transaction mode"
         );
 
         // Test session mode - prepared statements should be disabled
         config.config.general.pooler_mode = PoolerMode::Session;
         config.config.general.prepared_statements = PreparedStatements::Extended;
-        assert!(
-            !config.prepared_statements(),
+        assert_eq!(
+            config.prepared_statements(),
+            PreparedStatements::Disabled,
             "Prepared statements should be disabled in session mode"
         );
 
         // Test session mode with full prepared statements - should still be disabled
         config.config.general.pooler_mode = PoolerMode::Session;
         config.config.general.prepared_statements = PreparedStatements::Full;
-        assert!(
-            !config.prepared_statements(),
+        assert_eq!(
+            config.prepared_statements(),
+            PreparedStatements::Disabled,
             "Prepared statements should be disabled in session mode even when set to Full"
         );
 
         // Test transaction mode with disabled prepared statements - should remain disabled
         config.config.general.pooler_mode = PoolerMode::Transaction;
         config.config.general.prepared_statements = PreparedStatements::Disabled;
-        assert!(!config.prepared_statements(), "Prepared statements should remain disabled when explicitly set to Disabled in transaction mode");
+        assert_eq!(config.prepared_statements(), PreparedStatements::Disabled, "Prepared statements should remain disabled when explicitly set to Disabled in transaction mode");
     }
 
     #[test]

pgdog/src/config/pooling.rs

@@ -1,12 +1,13 @@
 use serde::{Deserialize, Serialize};
 use std::str::FromStr;
 
-#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)]
+#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Copy)]
 #[serde(rename_all = "snake_case")]
 pub enum PreparedStatements {
     Disabled,
     #[default]
     Extended,
+    ExtendedAnonymous,
     Full,
 }
 
@@ -18,6 +19,10 @@ impl PreparedStatements {
     pub fn enabled(&self) -> bool {
         !matches!(self, PreparedStatements::Disabled)
     }
+
+    pub fn rewrite_anonymous(&self) -> bool {
+        matches!(self, PreparedStatements::ExtendedAnonymous)
+    }
 }
 
 impl FromStr for PreparedStatements {
@@ -27,6 +32,7 @@ impl FromStr for PreparedStatements {
         match s.to_lowercase().as_str() {
             "disabled" => Ok(Self::Disabled),
             "extended" => Ok(Self::Extended),
+            "extended_anonymous" => Ok(Self::ExtendedAnonymous),
             "full" => Ok(Self::Full),
             _ => Err(format!("Invalid prepared statements mode: {}", s)),
         }

pgdog/src/frontend/client/mod.rs

@@ -482,7 +482,7 @@ impl Client {
         // Check config once per request.
         let config = config::config();
         // Configure prepared statements cache.
-        self.prepared_statements.enabled = config.prepared_statements();
+        self.prepared_statements.level = config.prepared_statements();
         self.prepared_statements.capacity = config.config.general.prepared_statements_limit;
         self.timeouts = Timeouts::from_config(&config.config.general);
 

pgdog/src/frontend/client/query_engine/prepared_statements.rs

@@ -1,3 +1,5 @@
+use crate::config::PreparedStatements;
+
 use super::*;
 
 impl QueryEngine {
@@ -7,8 +9,15 @@ impl QueryEngine {
         context: &mut QueryEngineContext<'_>,
     ) -> Result<(), Error> {
         for message in context.client_request.iter_mut() {
-            if message.extended() && context.prepared_statements.enabled {
-                context.prepared_statements.maybe_rewrite(message)?;
+            if message.extended() {
+                let level = context.prepared_statements.level;
+                match (level, message.anonymous()) {
+                    (PreparedStatements::ExtendedAnonymous, _)
+                    | (PreparedStatements::Extended, false) => {
+                        context.prepared_statements.maybe_rewrite(message)?
+                    }
+                    _ => (),
+                }
             }
         }
         Ok(())

pgdog/src/frontend/client/test/mod.rs

@@ -10,10 +10,10 @@ use bytes::{Buf, BufMut, BytesMut};
 
 use crate::{
     backend::databases::databases,
-    config::{config, load_test, load_test_replicas, set, Role},
+    config::{config, load_test, load_test_replicas, set, PreparedStatements, Role},
     frontend::{
         client::{BufferEvent, QueryEngine},
-        Client,
+        prepared_statements, Client,
     },
     net::{
         bind::Parameter, Bind, Close, CommandComplete, DataRow, Describe, ErrorResponse, Execute,
@@ -307,8 +307,8 @@ async fn test_client_with_replicas() {
             client.run().await.unwrap();
         });
         let buf = buffer!(
-            { Parse::new_anonymous("SELECT * FROM test_client_with_replicas") },
-            { Bind::new_statement("") },
+            { Parse::named("test", "SELECT * FROM test_client_with_replicas") },
+            { Bind::new_statement("test") },
             { Execute::new() },
             { Sync }
         );
@@ -769,3 +769,75 @@ async fn test_client_login_timeout_does_not_affect_queries() {
     conn.write_all(&buffer!({ Terminate })).await.unwrap();
     handle.await.unwrap();
 }
+
+#[tokio::test]
+async fn test_anon_prepared_statements() {
+    crate::logger();
+    load_test();
+
+    let (mut conn, mut client, _) = new_client!(false);
+
+    let mut c = (*config()).clone();
+    c.config.general.prepared_statements = PreparedStatements::ExtendedAnonymous;
+    set(c).unwrap();
+
+    let handle = tokio::spawn(async move {
+        client.run().await.unwrap();
+    });
+
+    conn.write_all(&buffer!(
+        { Parse::new_anonymous("SELECT 1") },
+        { Bind::new_params("", &[]) },
+        { Execute::new() },
+        { Sync }
+    ))
+    .await
+    .unwrap();
+
+    let _ = read!(conn, ['1', '2', 'D', 'C', 'Z']);
+
+    {
+        let cache = prepared_statements::PreparedStatements::global();
+        let read = cache.read();
+        assert!(!read.is_empty());
+    }
+
+    conn.write_all(&buffer!({ Terminate })).await.unwrap();
+    handle.await.unwrap();
+}
+
+#[tokio::test]
+async fn test_anon_prepared_statements_extended() {
+    crate::logger();
+    load_test();
+
+    let (mut conn, mut client, _) = new_client!(false);
+
+    let mut c = (*config()).clone();
+    c.config.general.prepared_statements = PreparedStatements::Extended;
+    set(c).unwrap();
+
+    let handle = tokio::spawn(async move {
+        client.run().await.unwrap();
+    });
+
+    conn.write_all(&buffer!(
+        { Parse::new_anonymous("SELECT 1") },
+        { Bind::new_params("", &[]) },
+        { Execute::new() },
+        { Sync }
+    ))
+    .await
+    .unwrap();
+
+    let _ = read!(conn, ['1', '2', 'D', 'C', 'Z']);
+
+    {
+        let cache = prepared_statements::PreparedStatements::global();
+        let read = cache.read();
+        assert!(read.is_empty());
+    }
+
+    conn.write_all(&buffer!({ Terminate })).await.unwrap();
+    handle.await.unwrap();
+}

pgdog/src/frontend/prepared_statements/mod.rs

@@ -6,6 +6,7 @@ use once_cell::sync::Lazy;
 use parking_lot::RwLock;
 
 use crate::{
+    config::PreparedStatements as PreparedStatementsLevel,
     frontend::router::parser::RewritePlan,
     net::{Parse, ProtocolMessage},
     stats::memory::MemoryUsage,
@@ -26,7 +27,7 @@ static CACHE: Lazy<PreparedStatements> = Lazy::new(PreparedStatements::default);
 pub struct PreparedStatements {
     pub(super) global: Arc<RwLock<GlobalCache>>,
     pub(super) local: HashMap<String, String>,
-    pub(super) enabled: bool,
+    pub(super) level: PreparedStatementsLevel,
     pub(super) capacity: usize,
     pub(super) memory_used: usize,
 }
@@ -35,7 +36,7 @@ impl MemoryUsage for PreparedStatements {
     #[inline]
     fn memory_usage(&self) -> usize {
         self.local.memory_usage()
-            + self.enabled.memory_usage()
+            + std::mem::size_of::<PreparedStatementsLevel>()
             + self.capacity.memory_usage()
             + std::mem::size_of::<Arc<RwLock<GlobalCache>>>()
     }
@@ -46,7 +47,7 @@ impl Default for PreparedStatements {
         Self {
             global: Arc::new(RwLock::new(GlobalCache::default())),
             local: HashMap::default(),
-            enabled: true,
+            level: PreparedStatementsLevel::Extended,
             capacity: usize::MAX,
             memory_used: 0,
         }

pgdog/src/net/protocol_message.rs

@@ -31,6 +31,17 @@ impl ProtocolMessage {
         )
     }
 
+    pub fn anonymous(&self) -> bool {
+        use ProtocolMessage::*;
+
+        match self {
+            Bind(bind) => bind.anonymous(),
+            Parse(parse) => parse.anonymous(),
+            Describe(describe) => describe.anonymous(),
+            _ => false,
+        }
+    }
+
     pub fn len(&self) -> usize {
         match self {
             Self::Bind(bind) => bind.len(),