fix: prepared statements with simple protocol not working (#628)

levkk · web-flow · commit 42da082ad95b · 2025-11-24T19:43:03.000-08:00
- fix: prepared statements with simple protocol not working #604
diff --git a/integration/prepared_statements_full/pgdog.toml b/integration/prepared_statements_full/pgdog.toml
@@ -0,0 +1,9 @@
+[general]
+prepared_statements = "full"
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+
+[admin]
+password = "pgdog"
diff --git a/integration/prepared_statements_full/users.toml b/integration/prepared_statements_full/users.toml
@@ -0,0 +1,4 @@
+[[users]]
+database = "pgdog"
+name = "pgdog"
+password = "pgdog"
diff --git a/integration/python/test_prepared.py b/integration/python/test_prepared.py
@@ -1,25 +1,33 @@
-from globals import normal_sync
-import pytest
+from globals import normal_sync, no_out_of_sync, admin
+from multiprocessing import Pool
+import time
 
 
-@pytest.mark.skip(reason="These are not working")
-def test_prepared_full():
-    for _ in range(5):
-        conn = normal_sync()
-        conn.autocommit = True
-
-        cur = conn.cursor()
-        cur.execute("PREPARE test_stmt AS SELECT 1")
-        cur.execute("PREPARE test_stmt AS SELECT 2")
-
+def run_prepare_execute(worker_id):
     conn = normal_sync()
     conn.autocommit = True
-    for _ in range(5):
-        cur = conn.cursor()
-
-        for i in range(5):
-            cur.execute(f"PREPARE test_stmt_{i} AS SELECT $1::bigint")
-            cur.execute(f"EXECUTE test_stmt_{i}({i})")
-            result = cur.fetchone()
-            assert result[0] == i
-        conn.commit()
+    cur = conn.cursor()
+
+    stmt_name = f"stmt_{worker_id}"
+    cur.execute(f"PREPARE {stmt_name} AS SELECT $1::bigint * 2")
+    time.sleep(0.01)
+
+    for i in range(100):
+        cur.execute(f"EXECUTE {stmt_name}({i})")
+        result = cur.fetchone()
+        assert result[0] == i * 2
+        time.sleep(0.01)
+
+    cur.execute(f"DEALLOCATE {stmt_name}")
+    conn.close()
+    return True
+
+
+def test_prepare_execute_parallel():
+    admin().execute("SET prepared_statements TO 'full'")
+
+    with Pool(5) as pool:
+        results = pool.map(run_prepare_execute, range(5))
+    assert all(results)
+    no_out_of_sync()
+    admin().execute("RELOAD")
diff --git a/pgdog/src/backend/pool/guard.rs b/pgdog/src/backend/pool/guard.rs
@@ -727,4 +727,53 @@ mod test {
         assert!(server.needs_drain());
         assert!(server.in_transaction());
     }
+
+    #[tokio::test]
+    async fn test_cleanup_syncs_prepared_statements() {
+        crate::logger();
+
+        let mut server = Guard::new(
+            Pool::new_test(),
+            Box::new(test_server().await),
+            Instant::now(),
+        );
+
+        assert!(!server.sync_prepared());
+
+        server
+            .send(&vec![Query::new("PREPARE test_stmt AS SELECT $1::bigint").into()].into())
+            .await
+            .unwrap();
+
+        for c in ['C', 'Z'] {
+            let msg = server.read().await.unwrap();
+            assert_eq!(msg.code(), c);
+        }
+
+        assert!(
+            server.sync_prepared(),
+            "sync_prepared flag should be set after PREPARE command"
+        );
+
+        let mut guard = server;
+        let mut server = guard.server.take().unwrap();
+        let cleanup = Cleanup::new(&guard, &mut server);
+
+        Guard::cleanup_internal(&mut server, cleanup, ConnectionRecovery::Recover)
+            .await
+            .unwrap();
+
+        assert!(
+            !server.sync_prepared(),
+            "sync_prepared flag should be cleared after cleanup"
+        );
+
+        assert!(
+            server.prepared_statements_mut().contains("test_stmt"),
+            "Statement should be in local cache after sync"
+        );
+
+        let one: Vec<i32> = server.fetch_all("SELECT 1").await.unwrap();
+        assert_eq!(one[0], 1);
+    }
 }
diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs
@@ -174,7 +174,16 @@ impl PreparedStatements {
                     self.state.add('3');
                 }
             }
-            ProtocolMessage::Prepare { .. } => (),
+            ProtocolMessage::Prepare { name, .. } => {
+                if self.contains(name) {
+                    return Ok(HandleResult::Drop);
+                } else {
+                    self.parses.push_back(name.clone());
+                    self.state.add_ignore('C');
+                    self.state.add_ignore('Z');
+                    return Ok(HandleResult::Forward);
+                }
+            }
             ProtocolMessage::CopyDone(_) => {
                 self.state.action('c')?;
             }
@@ -218,7 +227,7 @@ impl PreparedStatements {
                 self.describes.pop_front();
             }
 
-            '1' => {
+            '1' | 'C' => {
                 if let Some(name) = self.parses.pop_front() {
                     self.prepared(&name);
                 }
diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs
@@ -721,6 +721,7 @@ impl Server {
 
         let count = self.prepared_statements.len();
         self.stats.set_prepared_statements(count);
+        self.sync_prepared = false;
 
         Ok(())
     }
@@ -2304,4 +2305,40 @@ pub mod test {
             final_idle_time,
         );
     }
+
+    #[tokio::test]
+    async fn test_prepare_forces_sync_prepared_flag() {
+        let mut server = test_server().await;
+
+        assert!(!server.sync_prepared());
+
+        server
+            .send(&vec![Query::new("PREPARE test_stmt AS SELECT $1::bigint").into()].into())
+            .await
+            .unwrap();
+
+        for c in ['C', 'Z'] {
+            let msg = server.read().await.unwrap();
+            assert_eq!(msg.code(), c);
+        }
+
+        assert!(
+            server.sync_prepared(),
+            "sync_prepared flag should be set after PREPARE command"
+        );
+
+        server.sync_prepared_statements().await.unwrap();
+
+        assert!(
+            !server.sync_prepared(),
+            "sync_prepared flag should be cleared after sync_prepared_statements()"
+        );
+
+        server.execute("SELECT 1").await.unwrap();
+
+        assert!(
+            !server.sync_prepared(),
+            "sync_prepared flag should remain false after regular queries"
+        );
+    }
 }
diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs
@@ -48,7 +48,7 @@ impl QueryEngine {
         }
 
         if let Some(sql) = route.rewritten_sql() {
-            match context.client_request.rewrite(sql) {
+            match context.client_request.rewrite(&[Query::new(sql).into()]) {
                 Ok(()) => (),
                 Err(crate::net::Error::OnlySimpleForRewrites) => {
                     context.client_request.rewrite_prepared(
diff --git a/pgdog/src/frontend/client_request.rs b/pgdog/src/frontend/client_request.rs
@@ -7,7 +7,7 @@ use regex::Regex;
 use crate::{
     frontend::router::parser::RewritePlan,
     net::{
-        messages::{Bind, CopyData, Protocol, Query},
+        messages::{Bind, CopyData, Protocol},
         Error, Flush, ProtocolMessage,
     },
     stats::memory::MemoryUsage,
@@ -180,12 +180,12 @@ impl ClientRequest {
     }
 
     /// Rewrite query in buffer.
-    pub fn rewrite(&mut self, query: &str) -> Result<(), Error> {
+    pub fn rewrite(&mut self, request: &[ProtocolMessage]) -> Result<(), Error> {
         if self.messages.iter().any(|c| c.code() != 'Q') {
             return Err(Error::OnlySimpleForRewrites);
         }
         self.messages.clear();
-        self.messages.push(Query::new(query).into());
+        self.messages.extend(request.to_vec());
         Ok(())
     }
 
@@ -316,7 +316,7 @@ impl DerefMut for ClientRequest {
 
 #[cfg(test)]
 mod test {
-    use crate::net::{Describe, Execute, Parse, Sync};
+    use crate::net::{Describe, Execute, Parse, Query, Sync};
 
     use super::*;
 
diff --git a/pgdog/src/frontend/prepared_statements/mod.rs b/pgdog/src/frontend/prepared_statements/mod.rs
@@ -102,6 +102,8 @@ impl PreparedStatements {
         self.global.read().rewrite_plan(name)
     }
 
+    /// Set rewrite plan for UPDATE / INSERT statement
+    /// rewrites used in rewritten cross-shard queries.
     pub fn update_and_set_rewrite_plan(
         &mut self,
         name: &str,
@@ -117,6 +119,14 @@ impl PreparedStatements {
         self.local.get(name)
     }
 
+    /// Get globally-prepared statement by local name.
+    pub fn parse(&self, name: &str) -> Option<Parse> {
+        self.local
+            .get(name)
+            .map(|name| self.global.read().parse(name))
+            .flatten()
+    }
+
     /// Number of prepared statements in the local cache.
     pub fn len_local(&self) -> usize {
         self.local.len()
diff --git a/pgdog/src/frontend/router/parser/command.rs b/pgdog/src/frontend/router/parser/command.rs
@@ -1,7 +1,7 @@
 use super::*;
 use crate::{
     frontend::{client::TransactionType, BufferedQuery},
-    net::parameter::ParameterValue,
+    net::{parameter::ParameterValue, ProtocolMessage},
 };
 use lazy_static::lazy_static;
 
@@ -28,7 +28,7 @@ pub enum Command {
         value: ParameterValue,
     },
     PreparedStatement(Prepare),
-    Rewrite(String),
+    Rewrite(Vec<ProtocolMessage>),
     Shards(usize),
     Deallocate,
     Discard {
diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs
@@ -96,4 +96,7 @@ pub enum Error {
 
     #[error("cross-shard truncate not supported when schema-sharding is used")]
     CrossShardTruncateSchemaSharding,
+
+    #[error("prepared statement \"{0}\" doesn't exist")]
+    PreparedStatementDoesntExist(String),
 }
diff --git a/pgdog/src/frontend/router/parser/rewrite/mod.rs b/pgdog/src/frontend/router/parser/rewrite/mod.rs
@@ -5,8 +5,8 @@ use super::{Command, Error};
 mod insert_split;
 mod shard_key;
 
-use crate::frontend::PreparedStatements;
-use crate::net::Parse;
+use crate::net::{Parse, ProtocolMessage};
+use crate::{frontend::PreparedStatements, net::Query};
 pub use insert_split::{InsertSplitPlan, InsertSplitRow};
 pub use shard_key::{Assignment, AssignmentValue, ShardKeyRewritePlan};
 
@@ -46,17 +46,38 @@ impl<'a> Rewrite<'a> {
                 if let Some(ref mut node) = stmt.node {
                     match node {
                         NodeEnum::PrepareStmt(ref mut stmt) => {
-                            let statement = stmt.query.as_ref().ok_or(Error::EmptyQuery)?;
-                            let statement = statement.deparse().map_err(|_| Error::EmptyQuery)?;
+                            let statement = stmt
+                                .query
+                                .as_ref()
+                                .ok_or(Error::EmptyQuery)?
+                                .deparse()
+                                .map_err(|_| Error::EmptyQuery)?;
+
                             let mut parse = Parse::named(&stmt.name, &statement);
                             prepared_statements.insert_anyway(&mut parse);
                             stmt.name = parse.name().to_string();
+
+                            return Ok(Command::Rewrite(vec![Query::new(
+                                ast.deparse().map_err(|_| Error::EmptyQuery)?,
+                            )
+                            .into()]));
                         }
 
                         NodeEnum::ExecuteStmt(ref mut stmt) => {
-                            let name = prepared_statements.name(&stmt.name);
-                            if let Some(name) = name {
-                                stmt.name = name.to_string();
+                            let parse = prepared_statements.parse(&stmt.name);
+                            if let Some(parse) = parse {
+                                stmt.name = parse.name().to_string();
+
+                                return Ok(Command::Rewrite(vec![
+                                    ProtocolMessage::Prepare {
+                                        name: stmt.name.clone(),
+                                        statement: parse.query().to_string(),
+                                    },
+                                    Query::new(ast.deparse().map_err(|_| Error::EmptyQuery)?)
+                                        .into(),
+                                ]));
+                            } else {
+                                return Err(Error::PreparedStatementDoesntExist(stmt.name.clone()));
                             }
                         }
 
@@ -68,27 +89,30 @@ impl<'a> Rewrite<'a> {
             }
         }
 
-        Ok(Command::Rewrite(
-            ast.deparse().map_err(|_| Error::EmptyQuery)?,
-        ))
+        Err(Error::EmptyQuery)
     }
 }
 
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
 
+    use crate::net::{FromBytes, ToBytes};
+
     use super::*;
 
     #[test]
     fn test_rewrite_prepared() {
-        let ast = pg_query::parse("BEGIN; PREPARE test AS SELECT $1, $2, $3; PREPARE test2 AS SELECT * FROM my_table WHERE id = $1; COMMIT;").unwrap();
+        let ast = pg_query::parse("PREPARE test AS SELECT $1, $2, $3").unwrap();
         let rewrite = Rewrite::new(&ast);
         assert!(rewrite.needs_rewrite());
         let mut prepared_statements = PreparedStatements::new();
         let queries = rewrite.rewrite(&mut prepared_statements).unwrap();
         match queries {
-            Command::Rewrite(queries) => assert_eq!(queries, "BEGIN; PREPARE __pgdog_1 AS SELECT $1, $2, $3; PREPARE __pgdog_2 AS SELECT * FROM my_table WHERE id = $1; COMMIT"),
+            Command::Rewrite(messages) => {
+                let message = Query::from_bytes(messages[0].to_bytes().unwrap()).unwrap();
+                assert_eq!(message.query(), "PREPARE __pgdog_1 AS SELECT $1, $2, $3");
+            }
             _ => panic!("not a rewrite"),
         }
     }
diff --git a/pgdog/src/net/protocol_message.rs b/pgdog/src/net/protocol_message.rs
@@ -104,7 +104,9 @@ impl ToBytes for ProtocolMessage {
             Self::Bind(bind) => bind.to_bytes(),
             Self::Parse(parse) => parse.to_bytes(),
             Self::Describe(describe) => describe.to_bytes(),
-            Self::Prepare { statement, .. } => Query::new(statement).to_bytes(),
+            Self::Prepare { statement, name } => {
+                Query::new(format!("PREPARE {} AS {}", name, statement)).to_bytes()
+            }
             Self::Execute(execute) => execute.to_bytes(),
             Self::Close(close) => close.to_bytes(),
             Self::Query(query) => query.to_bytes(),

Original file line number	Diff line number	Diff line change
`@@ -174,7 +174,16 @@ impl PreparedStatements {`
`174`	`174`	`self.state.add('3');`
`175`	`175`	`}`
`176`	`176`	`}`
`177`		`- ProtocolMessage::Prepare { .. } => (),`
	`177`	`+ ProtocolMessage::Prepare { name, .. } => {`
	`178`	`+ if self.contains(name) {`
	`179`	`+ return Ok(HandleResult::Drop);`
	`180`	`+ } else {`
	`181`	`+ self.parses.push_back(name.clone());`
	`182`	`+ self.state.add_ignore('C');`
	`183`	`+ self.state.add_ignore('Z');`
	`184`	`+ return Ok(HandleResult::Forward);`
	`185`	`+ }`
	`186`	`+ }`
`178`	`187`	`ProtocolMessage::CopyDone(_) => {`
`179`	`188`	`self.state.action('c')?;`
`180`	`189`	`}`
`@@ -218,7 +227,7 @@ impl PreparedStatements {`
`218`	`227`	`self.describes.pop_front();`
`219`	`228`	`}`
`220`	`229`
`221`		`- '1' => {`
	`230`	`+ '1' \| 'C' => {`
`222`	`231`	`if let Some(name) = self.parses.pop_front() {`
`223`	`232`	`self.prepared(&name);`
`224`	`233`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ impl QueryEngine {`
`48`	`48`	`}`
`49`	`49`
`50`	`50`	`if let Some(sql) = route.rewritten_sql() {`
`51`		`- match context.client_request.rewrite(sql) {`
	`51`	`+ match context.client_request.rewrite(&[Query::new(sql).into()]) {`
`52`	`52`	`Ok(()) => (),`
`53`	`53`	`Err(crate::net::Error::OnlySimpleForRewrites) => {`
`54`	`54`	`context.client_request.rewrite_prepared(`