feat: rewrite simple to prepared

.schema/pgdog.schema.json

@@ -220,6 +220,7 @@
         "enabled": false,
         "primary_key": "ignore",
         "shard_key": "error",
+        "simple_to_prepared": false,
         "split_inserts": "error"
       }
     },
@@ -1870,6 +1871,11 @@
           "$ref": "#/$defs/RewriteMode",
           "default": "error"
         },
+        "simple_to_prepared": {
+          "description": "Rewrite simple queries to prepared statements.",
+          "type": "boolean",
+          "default": false
+        },
         "split_inserts": {
           "description": "Behavior for multi-row `INSERT` on sharded tables: `error` rejects, `rewrite` distributes rows to their shards, `ignore` forwards unchanged.\n\n_Default:_ `error`\n\n<https://docs.pgdog.dev/configuration/pgdog.toml/rewrite/#split_inserts>",
           "$ref": "#/$defs/RewriteMode",

cli.sh

@@ -17,7 +17,7 @@ function admin() {
 # - protocol: simple|extended|prepared
 #
 function bench() {
-    PGPASSWORD=pgdog pgbench -h 127.0.0.1 -p 6432 -U pgdog pgdog --protocol ${1:-simple} -t 100000000 -c 10 -P 1 -f pgdog/tests/pgbouncer/pgbench-parser.sql
+    PGPASSWORD=pgdog pgbench -h 127.0.0.1 -p 6432 -U pgdog pgdog --protocol ${1:-simple} -t 100000000 -c 10 -P 1 -S
 }
 
 function bench_init() {

integration/simple_to_prepared/pgdog.toml

@@ -0,0 +1,11 @@
+[general]
+idle_healthcheck_delay = 1000000000
+query_parser = "on"
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+
+
+[rewrite]
+simple_to_prepared = false

integration/simple_to_prepared/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+name = "pgdog"
+password = "pgdog"
+database = "pgdog"

pgdog-config/src/rewrite.rs

@@ -88,6 +88,10 @@ pub struct Rewrite {
     /// <https://docs.pgdog.dev/configuration/pgdog.toml/rewrite/#primary_key>
     #[serde(default = "Rewrite::default_primary_key")]
     pub primary_key: RewriteMode,
+
+    /// Rewrite simple queries to prepared statements.
+    #[serde(default)]
+    pub simple_to_prepared: bool,
 }
 
 impl Default for Rewrite {
@@ -97,6 +101,7 @@ impl Default for Rewrite {
             shard_key: Self::default_shard_key(),
             split_inserts: Self::default_split_inserts(),
             primary_key: Self::default_primary_key(),
+            simple_to_prepared: bool::default(),
         }
     }
 }

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

@@ -242,12 +242,20 @@ impl QueryEngine {
         // Do this before flushing, because flushing can take time.
         self.cleanup_backend(context)?;
 
-        trace!("{:#?} >>> {:?}", message, context.stream.peer_addr());
-
-        if flush {
-            context.stream.send_flush(&message).await?;
-        } else {
-            context.stream.send(&message).await?;
+        let forward_to_client = context
+            .rewrite_result
+            .as_ref()
+            .map(|rewrite| rewrite.apply_after_execution(&message).forward())
+            .unwrap_or(true);
+
+        if forward_to_client {
+            trace!("{:#?} >>> {:?}", message, context.stream.peer_addr());
+
+            if flush {
+                context.stream.send_flush(&message).await?;
+            } else {
+                context.stream.send(&message).await?;
+            }
         }
 
         if code == 'Z' {

pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs

@@ -18,7 +18,7 @@ async fn run_test(messages: Vec<ProtocolMessage>) -> Option<OffsetPlan> {
     engine.parse_and_rewrite(&mut context).await.unwrap();
 
     match context.rewrite_result {
-        Some(RewriteResult::InPlace { offset }) => offset,
+        Some(RewriteResult::InPlace { offset, .. }) => offset,
         other => panic!("expected InPlace, got {:?}", other),
     }
 }

pgdog/src/frontend/prepared_statements/mod.rs

@@ -2,6 +2,7 @@
 
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
+use bytes::Bytes;
 use once_cell::sync::Lazy;
 use parking_lot::RwLock;
 use tracing::debug;
@@ -34,6 +35,7 @@ pub struct PreparedStatements {
     // mapping the client statement name -> __pgdog__ name from global cache
     pub(super) local: HashMap<String, String>,
     pub(super) level: PreparedStatementsLevel,
+    rewritten_simple_to_prepared: HashMap<Bytes, String>,
     pub(super) memory_used: usize,
 }
 
@@ -43,6 +45,7 @@ impl Default for PreparedStatements {
             global: Arc::new(RwLock::new(GlobalCache::default())),
             local: HashMap::default(),
             level: PreparedStatementsLevel::Extended,
+            rewritten_simple_to_prepared: HashMap::new(),
             memory_used: 0,
         }
     }
@@ -66,8 +69,28 @@ impl PreparedStatements {
         Ok(())
     }
 
+    /// Manually map a local prepared statement to a global one.
+    ///
+    /// Warning: don't use this unless you understand the side-effects:
+    ///
+    /// 1. When client disconnects, this statement's global counter will be decreased by 1.
+    /// 2. The statement will not be removed from the global cache until the client disconnects
+    ///    because clients are not aware of this and will never close it.
+    ///
+    pub(crate) fn insert_rewritten_simple_to_prepared(&mut self, parse: &Parse) -> String {
+        if let Some(name) = self.rewritten_simple_to_prepared.get(&parse.query_ref()) {
+            name.to_owned()
+        } else {
+            let (_new, name) = { self.global.write().insert(parse) };
+            self.local.insert(name.to_owned(), name.to_owned());
+            self.rewritten_simple_to_prepared
+                .insert(parse.query_ref(), name.clone());
+            name
+        }
+    }
+
     /// Register prepared statement with the global cache.
-    pub fn insert(&mut self, parse: &mut Parse) {
+    pub(crate) fn insert(&mut self, parse: &mut Parse) {
         let (_new, name) = { self.global.write().insert(parse) };
         let key = parse.name();
         let existed = self.local.insert(key.to_owned(), name.clone());

pgdog/src/frontend/router/parser/cache/ast.rs

@@ -68,7 +68,7 @@ impl Deref for Ast {
 
 impl Ast {
     /// Parse statement and run the rewrite engine, if necessary.
-    pub(super) fn new(
+    fn new(
         query: &AstQuery,
         schema: &ShardingSchema,
         db_schema: &Schema,
@@ -78,6 +78,7 @@ impl Ast {
     ) -> Result<Self, Error> {
         let now = Instant::now();
         let ast = pg_raw_parse::parse(query.query_without_comment).map_err(Error::Parse)?;
+        let multiple_statements = ast.stmts().count() > 1;
 
         // Run the rewrite unconditionally. Even when a shard comment will
         // route the query to a specific shard, we need to know whether the
@@ -91,6 +92,7 @@ impl Ast {
             db_schema,
             user,
             search_path,
+            multiple_statements,
         });
         let mut rewrite_plan = Default::default();
         let ast = make::try_owned(|mem| {

pgdog/src/frontend/router/parser/cache/test.rs

@@ -359,3 +359,28 @@ fn test_truncated_query_non_ascii_char_boundary() {
     let ast_query = AstQuery::from_query(&buffered);
     assert_eq!(ast_query.truncated_query(9), "SELECT '€");
 }
+
+#[test]
+fn rejects_rewritten_multi_statement_queries() {
+    let mut ctx = test_context();
+    ctx.sharding_schema.rewrite.simple_to_prepared = true;
+    let mut prepared_statements = PreparedStatements::default();
+
+    let unchanged =
+        BufferedQuery::Query(Query::new("SELECT current_user; SELECT current_database()"));
+    Cache::get()
+        .query(&unchanged, &ctx, &mut prepared_statements)
+        .expect("multi-statement query without rewrites should parse");
+
+    let rewritten = BufferedQuery::Query(Query::new("SELECT 1; SELECT 2"));
+    let error = Cache::get()
+        .query(&rewritten, &ctx, &mut prepared_statements)
+        .expect_err("rewritten multi-statement query should be rejected");
+
+    assert!(matches!(
+        error,
+        crate::frontend::router::parser::Error::Rewrite(
+            crate::frontend::router::parser::rewrite::statement::Error::MultiStatementRewrite
+        )
+    ));
+}

pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs

@@ -500,6 +500,7 @@ mod tests {
             db_schema,
             user: "",
             search_path: None,
+            multiple_statements: false,
         });
         let mut plan = Default::default();
         let ast = make::try_owned(|mem| {

pgdog/src/frontend/router/parser/rewrite/statement/error.rs

@@ -40,4 +40,10 @@ pub enum Error {
 
     #[error("prepared statement '{0}' does not exist")]
     ExecuteMissingPrepare(String),
+
+    #[error("prepared statement: {0}")]
+    PreparedStmt(#[from] crate::frontend::prepared_statements::Error),
+
+    #[error("cannot rewrite a multi-statement query")]
+    MultiStatementRewrite,
 }

pgdog/src/frontend/router/parser/rewrite/statement/insert.rs

@@ -273,6 +273,7 @@ mod tests {
             db_schema: &db_schema,
             user: "",
             search_path: None,
+            multiple_statements: false,
         });
         let mut plan = RewritePlan::default();
         rewriter.split_insert(insert, &mut plan).unwrap();

pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

@@ -15,13 +15,15 @@ pub mod insert;
 pub mod offset;
 pub mod plan;
 pub mod simple_prepared;
+pub mod simple_to_prepared;
 pub mod unique_id;
 pub mod update;
 
 pub use error::Error;
 pub use insert::InsertSplit;
 pub(crate) use plan::RewritePlan;
 pub use simple_prepared::SimplePreparedResult;
+pub(crate) use simple_to_prepared::*;
 pub(crate) use update::*;
 
 /// Statement rewrite engine context.
@@ -42,6 +44,8 @@ pub struct StatementRewriteContext<'a> {
     pub user: &'a str,
     /// Search path for table lookups.
     pub search_path: Option<&'a ParameterValue>,
+    /// Whether the query contains more than one SQL statement.
+    pub multiple_statements: bool,
 }
 
 #[derive(Debug)]
@@ -65,6 +69,8 @@ pub struct StatementRewrite<'a> {
     user: &'a str,
     /// Search path for table lookups.
     search_path: Option<&'a ParameterValue>,
+    /// Whether the query contains more than one SQL statement.
+    multiple_statements: bool,
 }
 
 impl<'a> StatementRewrite<'a> {
@@ -82,6 +88,7 @@ impl<'a> StatementRewrite<'a> {
             db_schema: ctx.db_schema,
             user: ctx.user,
             search_path: ctx.search_path,
+            multiple_statements: ctx.multiple_statements,
         }
     }
 
@@ -104,6 +111,10 @@ impl<'a> StatementRewrite<'a> {
     ) -> Result<RewritePlan, Error> {
         let mut plan = RewritePlan::default();
 
+        // N.B. The simple to prepared rewriter should run first.
+        // All subsequent rewriters will act on the prepared statement.
+        self.rewrite_simple_to_prepared(stmt.stmt_mut(), mem, &mut plan)?;
+
         match stmt.stmt() {
             Node::InsertStmt(_)
             | Node::SelectStmt(_)
@@ -162,8 +173,19 @@ impl<'a> StatementRewrite<'a> {
             self.limit_offset(&select, &mut plan);
         }
 
+        if self.rewritten && self.multiple_statements {
+            return Err(Error::MultiStatementRewrite);
+        }
+
         if self.rewritten {
-            plan.stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned());
+            let stmt = pg_raw_parse::deparse(&*stmt)?.as_str().to_owned();
+
+            // N.B. careful with ordering. This should run before insert splits, etc.
+            // since we want to make sure the statement is registered with the global cache.
+            plan.simple_to_prepared
+                .step_two(self.prepared_statements, &stmt)?;
+
+            plan.stmt = Some(stmt);
         }
 
         if let Node::InsertStmt(insert) = stmt.stmt() {

pgdog/src/frontend/router/parser/rewrite/statement/offset.rs

@@ -248,6 +248,7 @@ mod tests {
             db_schema: &db_schema,
             user: "test",
             search_path: None,
+            multiple_statements: false,
         });
         let mut plan = RewritePlan::default();
         rewrite.limit_offset(