From 7486149643a04b8aea428acaa21b8be0d871769c Mon Sep 17 00:00:00 2001
From: PThorpe92 <preston@unlockedlabs.org>
Date: Wed, 26 Mar 2025 13:22:12 -0400
Subject: [PATCH] Support LIMIT clause on update queries

---
 core/translate/emitter.rs | 17 ++++++++++-------
 core/translate/update.rs  |  2 +-
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs
index b2cea311f..ab27d2715 100644
--- a/core/translate/emitter.rs
+++ b/core/translate/emitter.rs
@@ -437,6 +437,15 @@ fn emit_program_for_update(
         &plan.table_references,
         OperationMode::UPDATE,
     )?;
+    if t_ctx.reg_limit.is_none() && plan.limit.is_some() {
+        let reg = program.alloc_register();
+        t_ctx.reg_limit = Some(reg);
+        program.emit_insn(Insn::Integer {
+            value: plan.limit.unwrap() as i64,
+            dest: reg,
+        });
+        program.mark_last_insn_constant();
+    }
     open_loop(
         program,
         &mut t_ctx,
@@ -553,13 +562,7 @@ fn emit_update_insns(
     });
     program.emit_insn(Insn::InsertAwait { cursor_id });
 
-    if let Some(limit) = plan.limit {
-        let limit_reg = program.alloc_register();
-        program.emit_insn(Insn::Integer {
-            value: limit as i64,
-            dest: limit_reg,
-        });
-        program.mark_last_insn_constant();
+    if let Some(limit_reg) = t_ctx.reg_limit {
         program.emit_insn(Insn::DecrJumpZero {
             reg: limit_reg,
             target_pc: t_ctx.label_main_loop_end.unwrap(),
diff --git a/core/translate/update.rs b/core/translate/update.rs
index e552777d6..086e7719b 100644
--- a/core/translate/update.rs
+++ b/core/translate/update.rs
@@ -155,7 +155,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result<
         Some(&result_columns),
         &mut where_clause,
     )?;
-    let limit = if let Some(Ok((_, limit))) = body.limit.as_ref().map(|l| parse_limit(*l.clone())) {
+    let limit = if let Some(Ok((limit, _))) = body.limit.as_ref().map(|l| parse_limit(*l.clone())) {
         limit
     } else {
         None