Support LIMIT clause on update queries

This commit is contained in:
PThorpe92
2025-03-26 13:22:12 -04:00
parent b7fca31ef6
commit 7486149643
2 changed files with 11 additions and 8 deletions

View File

@@ -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(),

View File

@@ -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