From f8eb4ba14d67b63bb9383dda395b186c8664e838 Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Mon, 21 Jul 2025 20:51:42 -0300
Subject: [PATCH] implement reprepare for statements

---
 core/lib.rs      | 49 ++++++++++++++++++++++++++++++++++++++++++++++--
 core/vdbe/mod.rs |  4 ++--
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/core/lib.rs b/core/lib.rs
index 6d15c9948..3dc0a5a9b 100644
--- a/core/lib.rs
+++ b/core/lib.rs
@@ -1743,8 +1743,53 @@ impl Statement {
     }
 
     pub fn step(&mut self) -> Result<StepResult> {
-        self.program
-            .step(&mut self.state, self.mv_store.clone(), self.pager.clone())
+        const MAX_SCHEMA_RETRY: usize = 50;
+        let mut res = self
+            .program
+            .step(&mut self.state, self.mv_store.clone(), self.pager.clone());
+        for _ in 0..MAX_SCHEMA_RETRY {
+            // Only reprepare if we still need to update schema
+            if !matches!(res, Err(LimboError::SchemaUpdated)) {
+                break;
+            }
+            self.reprepare()?;
+            res = self
+                .program
+                .step(&mut self.state, self.mv_store.clone(), self.pager.clone());
+        }
+
+        res
+    }
+
+    fn reprepare(&mut self) -> Result<()> {
+        self.program = {
+            let mut parser = Parser::new(self.program.sql.as_bytes());
+            let cmd = parser.next()?;
+            let cmd = cmd.expect("Same SQL string should be able to be parsed");
+
+            let conn = self.program.connection.clone();
+            let syms = conn.syms.borrow();
+
+            match cmd {
+                Cmd::Stmt(stmt) => translate::translate(
+                    conn.schema.borrow().deref(),
+                    stmt,
+                    self.pager.clone(),
+                    conn.clone(),
+                    &syms,
+                    QueryMode::Normal,
+                    &self.program.sql,
+                )?,
+                Cmd::Explain(_stmt) => todo!(),
+                Cmd::ExplainQueryPlan(_stmt) => todo!(),
+            }
+        };
+        // Save parameters before they are reset
+        let parameters = std::mem::take(&mut self.state.parameters);
+        self.reset();
+        // Load the parameters back into the state
+        self.state.parameters = parameters;
+        Ok(())
     }
 
     pub fn run_once(&self) -> Result<()> {
diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs
index ed96043a1..7766b2bfd 100644
--- a/core/vdbe/mod.rs
+++ b/core/vdbe/mod.rs
@@ -248,7 +248,7 @@ pub struct ProgramState {
     regex_cache: RegexCache,
     pub(crate) mv_tx_id: Option<crate::mvcc::database::TxID>,
     interrupted: bool,
-    parameters: HashMap<NonZero<usize>, Value>,
+    pub parameters: HashMap<NonZero<usize>, Value>,
     commit_state: CommitState,
     #[cfg(feature = "json")]
     json_cache: JsonCacheCell,
@@ -384,7 +384,7 @@ pub struct Program {
     pub change_cnt_on: bool,
     pub result_columns: Vec<ResultSetColumn>,
     pub table_references: TableReferences,
-    pub sql: String
+    pub sql: String,
 }
 
 impl Program {