Merge 'translate/update: allow for updating nonalias'd rowid column explicitly ' from Preston Thorpe

closes #3276

Reviewed-by: Jussi Saurio <jussi.saurio@gmail.com>

Closes #3349
This commit is contained in:
Jussi Saurio
2025-09-26 09:13:03 +03:00
committed by GitHub
4 changed files with 89 additions and 28 deletions

View File

@@ -45,6 +45,9 @@ const SCHEMA_TABLE_NAME: &str = "sqlite_schema";
const SCHEMA_TABLE_NAME_ALT: &str = "sqlite_master";
pub const DBSP_TABLE_PREFIX: &str = "__turso_internal_dbsp_state_";
/// Used to refer to the implicit rowid column in tables without an alias during UPDATE
pub const ROWID_SENTINEL: usize = usize::MAX;
/// Check if a table name refers to a system table that should be protected from direct writes
pub fn is_system_table(table_name: &str) -> bool {
let normalized = table_name.to_lowercase();

View File

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

View File

@@ -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<usize> = 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()

View File

@@ -378,3 +378,11 @@ do_execsql_test_on_specific_db {:memory:} rowid-update-updates-all-indexes {
} {2|3|1
2|3|1
2|3|1}
# https://github.com/tursodatabase/turso/issues/3276
do_execsql_test_on_specific_db {:memory:} can-update-rowid-directly {
CREATE TABLE test (name TEXT);
INSERT INTO test (name) VALUES ('test');
UPDATE test SET rowid = 5;
SELECT rowid, name from test;
} {5|test}