diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs index fb5f35a24..48cb4a45e 100644 --- a/core/translate/emitter.rs +++ b/core/translate/emitter.rs @@ -23,7 +23,7 @@ use super::select::emit_simple_count; use super::subquery::emit_subqueries; use crate::error::SQLITE_CONSTRAINT_PRIMARYKEY; use crate::function::Func; -use crate::schema::{BTreeTable, Column, Schema, Table}; +use crate::schema::{BTreeTable, Column, Schema, Table, ROWID_SENTINEL}; use crate::translate::compound_select::emit_program_for_compound_select; use crate::translate::expr::{ emit_returning_results, translate_expr_no_constant_opt, walk_expr_mut, NoConstantOptReason, @@ -895,12 +895,16 @@ fn emit_update_insns( .iter() .position(|c| c.is_rowid_alias); + let has_direct_rowid_update = plan + .set_clauses + .iter() + .any(|(idx, _)| *idx == ROWID_SENTINEL); + let has_user_provided_rowid = if let Some(index) = rowid_alias_index { - plan.set_clauses.iter().position(|(idx, _)| *idx == index) + plan.set_clauses.iter().any(|(idx, _)| *idx == index) } else { - None - } - .is_some(); + has_direct_rowid_update + }; let rowid_set_clause_reg = if has_user_provided_rowid { Some(program.alloc_register()) @@ -958,9 +962,29 @@ fn emit_update_insns( let table_name = unsafe { &*table_ref }.table.get_name(); let start = if is_virtual { beg + 2 } else { beg + 1 }; + + if has_direct_rowid_update { + if let Some((_, expr)) = plan.set_clauses.iter().find(|(i, _)| *i == ROWID_SENTINEL) { + let rowid_set_clause_reg = rowid_set_clause_reg.unwrap(); + translate_expr( + program, + Some(&plan.table_references), + expr, + rowid_set_clause_reg, + &t_ctx.resolver, + )?; + program.emit_insn(Insn::MustBeInt { + reg: rowid_set_clause_reg, + }); + } + } for (idx, table_column) in unsafe { &*table_ref }.columns().iter().enumerate() { let target_reg = start + idx; - if let Some((_, expr)) = plan.set_clauses.iter().find(|(i, _)| *i == idx) { + if let Some((col_idx, expr)) = plan.set_clauses.iter().find(|(i, _)| *i == idx) { + // Skip if this is the sentinel value + if *col_idx == ROWID_SENTINEL { + continue; + } if has_user_provided_rowid && (table_column.primary_key || table_column.is_rowid_alias) && !is_virtual @@ -1298,8 +1322,8 @@ fn emit_update_insns( if has_user_provided_rowid { let record_label = program.allocate_label(); - let idx = rowid_alias_index.unwrap(); let target_reg = rowid_set_clause_reg.unwrap(); + program.emit_insn(Insn::Eq { lhs: target_reg, rhs: beg, @@ -1314,19 +1338,23 @@ fn emit_update_insns( target_pc: record_label, }); - program.emit_insn(Insn::Halt { - err_code: SQLITE_CONSTRAINT_PRIMARYKEY, - description: format!( - "{}.{}", - table_name, - unsafe { &*table_ref } + let description = if let Some(idx) = rowid_alias_index { + String::from(table_name) + + "." + + unsafe { &*table_ref } .columns() .get(idx) .unwrap() .name .as_ref() .map_or("", |v| v) - ), + } else { + String::from(table_name) + ".rowid" + }; + + program.emit_insn(Insn::Halt { + err_code: SQLITE_CONSTRAINT_PRIMARYKEY, + description, }); program.preassign_label_to_next_insn(record_label); diff --git a/core/translate/update.rs b/core/translate/update.rs index 447b3c8ee..d4c4c1fbd 100644 --- a/core/translate/update.rs +++ b/core/translate/update.rs @@ -1,13 +1,13 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use crate::schema::{BTreeTable, Column, Type}; +use crate::schema::{BTreeTable, Column, Type, ROWID_SENTINEL}; use crate::translate::expr::{ bind_and_rewrite_expr, walk_expr, BindingBehavior, ParamState, WalkControl, }; use crate::translate::optimizer::optimize_select_plan; use crate::translate::plan::{Operation, QueryDestination, Scan, Search, SelectPlan}; -use crate::translate::planner::parse_limit; +use crate::translate::planner::{parse_limit, ROWID_STRS}; use crate::vdbe::builder::CursorType; use crate::{ bail_parse_error, @@ -214,18 +214,39 @@ pub fn prepare_update_plan( ); } - // Map each column to its corresponding expression for (col_name, expr) in set.col_names.iter().zip(values.iter()) { let ident = normalize_ident(col_name.as_str()); - let col_index = match column_lookup.get(&ident) { - Some(idx) => idx, - None => bail_parse_error!("no such column: {}", ident), - }; - // Update existing entry or add new one - match set_clauses.iter_mut().find(|(idx, _)| idx == col_index) { - Some((_, existing_expr)) => *existing_expr = expr.clone(), - None => set_clauses.push((*col_index, expr.clone())), + // Check if this is the 'rowid' keyword + if ROWID_STRS.iter().any(|s| s.eq_ignore_ascii_case(&ident)) { + // Find the rowid alias column if it exists + if let Some((idx, _col)) = table + .columns() + .iter() + .enumerate() + .find(|(_, c)| c.is_rowid_alias) + { + // Use the rowid alias column index + match set_clauses.iter_mut().find(|(i, _)| i == &idx) { + Some((_, existing_expr)) => *existing_expr = expr.clone(), + None => set_clauses.push((idx, expr.clone())), + } + } else { + // No rowid alias, use sentinel value for actual rowid + match set_clauses.iter_mut().find(|(i, _)| *i == ROWID_SENTINEL) { + Some((_, existing_expr)) => *existing_expr = expr.clone(), + None => set_clauses.push((ROWID_SENTINEL, expr.clone())), + } + } + } else { + let col_index = match column_lookup.get(&ident) { + Some(idx) => idx, + None => bail_parse_error!("no such column: {}", ident), + }; + match set_clauses.iter_mut().find(|(idx, _)| idx == col_index) { + Some((_, existing_expr)) => *existing_expr = expr.clone(), + None => set_clauses.push((*col_index, expr.clone())), + } } } } @@ -262,10 +283,11 @@ pub fn prepare_update_plan( let columns = table.columns(); let rowid_alias_used = set_clauses.iter().fold(false, |accum, (idx, _)| { - accum || columns[*idx].is_rowid_alias + accum || (*idx != ROWID_SENTINEL && columns[*idx].is_rowid_alias) }); + let direct_rowid_update = set_clauses.iter().any(|(idx, _)| *idx == ROWID_SENTINEL); - let (ephemeral_plan, mut where_clause) = if rowid_alias_used { + let (ephemeral_plan, mut where_clause) = if rowid_alias_used || direct_rowid_update { let mut where_clause = vec![]; let internal_id = program.table_reference_counter.next(); @@ -384,7 +406,7 @@ pub fn prepare_update_plan( let updated_cols: HashSet = set_clauses.iter().map(|(i, _)| *i).collect(); let rowid_alias_used = set_clauses .iter() - .any(|(idx, _)| columns[*idx].is_rowid_alias); + .any(|(idx, _)| *idx == ROWID_SENTINEL || columns[*idx].is_rowid_alias); let indexes_to_update = if rowid_alias_used { // If the rowid alias is used in the SET clause, we need to update all indexes indexes.cloned().collect()