From d04b07b8b74a28b72d55689c458616e353dbd8b8 Mon Sep 17 00:00:00 2001 From: PThorpe92 Date: Sat, 27 Sep 2025 20:48:42 -0400 Subject: [PATCH] Add pragma foreign_keys and fk_if_zero and fk_counter opcodes --- core/pragma.rs | 4 +++ core/translate/pragma.rs | 25 ++++++++++++++++- core/vdbe/execute.rs | 59 ++++++++++++++++++++++++++++++++++++++++ core/vdbe/explain.rs | 20 +++++++++++++- core/vdbe/insn.rs | 16 +++++++++++ core/vdbe/mod.rs | 2 ++ 6 files changed, 124 insertions(+), 2 deletions(-) diff --git a/core/pragma.rs b/core/pragma.rs index c238134e4..8cf9a99c5 100644 --- a/core/pragma.rs +++ b/core/pragma.rs @@ -131,6 +131,10 @@ pub fn pragma_for(pragma: &PragmaName) -> Pragma { PragmaFlags::NoColumns1 | PragmaFlags::Result0, &["mvcc_checkpoint_threshold"], ), + ForeignKeys => Pragma::new( + PragmaFlags::NoColumns1 | PragmaFlags::Result0, + &["foreign_keys"], + ), } } diff --git a/core/translate/pragma.rs b/core/translate/pragma.rs index 19542adad..f08bd5a15 100644 --- a/core/translate/pragma.rs +++ b/core/translate/pragma.rs @@ -4,7 +4,7 @@ use chrono::Datelike; use std::sync::Arc; use turso_macros::match_ignore_ascii_case; -use turso_parser::ast::{self, ColumnDefinition, Expr, Literal}; +use turso_parser::ast::{self, ColumnDefinition, Expr, Literal, Name}; use turso_parser::ast::{PragmaName, QualifiedName}; use super::integrity_check::translate_integrity_check; @@ -387,6 +387,21 @@ fn update_pragma( connection.set_mvcc_checkpoint_threshold(threshold)?; Ok((program, TransactionMode::None)) } + PragmaName::ForeignKeys => { + let enabled = match &value { + Expr::Literal(Literal::Keyword(name)) | Expr::Id(name) => { + let name_bytes = name.as_bytes(); + match_ignore_ascii_case!(match name_bytes { + b"ON" | b"TRUE" | b"YES" | b"1" => true, + _ => false, + }) + } + Expr::Literal(Literal::Numeric(n)) => !matches!(n.as_str(), "0"), + _ => false, + }; + connection.set_foreign_keys(enabled); + Ok((program, TransactionMode::None)) + } } } @@ -704,6 +719,14 @@ fn query_pragma( program.add_pragma_result_column(pragma.to_string()); Ok((program, TransactionMode::None)) } + PragmaName::ForeignKeys => { + let enabled = connection.foreign_keys_enabled(); + let register = program.alloc_register(); + program.emit_int(enabled as i64, register); + program.emit_result_row(register, 1); + program.add_pragma_result_column(pragma.to_string()); + Ok((program, TransactionMode::None)) + } } } diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 9a03ca9bc..4c93ded2b 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -8276,6 +8276,65 @@ fn handle_text_sum(acc: &mut Value, sum_state: &mut SumAggState, parsed_number: } } +pub fn op_fk_counter( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + pager: &Arc, + mv_store: Option<&Arc>, +) -> Result { + load_insn!( + FkCounter { + increment_value, + check_abort, + }, + insn + ); + state.fk_constraint_counter = state.fk_constraint_counter.saturating_add(*increment_value); + + // If check_abort is true and counter is negative, abort with constraint error + // This shouldn't happen in well-formed bytecode but acts as a safety check + if *check_abort && state.fk_constraint_counter < 0 { + return Err(LimboError::Constraint( + "FOREIGN KEY constraint failed".into(), + )); + } + + state.pc += 1; + Ok(InsnFunctionStepResult::Step) +} + +pub fn op_fk_if_zero( + program: &Program, + state: &mut ProgramState, + insn: &Insn, + _pager: &Arc, + _mv_store: Option<&Arc>, +) -> Result { + load_insn!(FkIfZero { target_pc, if_zero }, insn); + let fk_enabled = program.connection.foreign_keys_enabled(); + + // Jump if any: + // Foreign keys are disabled globally + // p1 is true AND deferred constraint counter is zero + // p1 is false AND deferred constraint counter is non-zero + let should_jump = if !fk_enabled { + true + } else if *if_zero { + state.fk_constraint_counter == 0 + } else { + state.fk_constraint_counter != 0 + }; + + if should_jump { + state.pc = target_pc.as_offset_int(); + } else { + state.pc += 1; + } + + Ok(InsnFunctionStepResult::Step) +} + mod cmath { extern "C" { pub fn exp(x: f64) -> f64; diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index 5e8dde2fe..15485bab7 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -1804,7 +1804,25 @@ pub fn insn_to_row( 0, String::new(), ), - } + Insn::FkCounter{check_abort, increment_value} => ( + "FkCounter", + *check_abort as i32, + *increment_value as i32, + 0, + Value::build_text(""), + 0, + String::new(), + ), + Insn::FkIfZero{target_pc, if_zero } => ( + "FkIfZero", + target_pc.as_debug_int(), + *if_zero as i32, + 0, + Value::build_text(""), + 0, + String::new(), + ), + } } pub fn insn_to_row_with_comment( diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index 67e1b784d..06e392902 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -1169,6 +1169,20 @@ pub enum Insn { p2: Option, // P2: address of parent explain instruction detail: String, // P4: detail text }, + // Increment a "constraint counter" by P2 (P2 may be negative or positive). + // If P1 is non-zero, the database constraint counter is incremented (deferred foreign key constraints). + // Otherwise, if P1 is zero, the statement counter is incremented (immediate foreign key constraints). + FkCounter { + check_abort: bool, + increment_value: isize, + }, + // This opcode tests if a foreign key constraint-counter is currently zero. If so, jump to instruction P2. Otherwise, fall through to the next instruction. + // If P1 is non-zero, then the jump is taken if the database constraint-counter is zero (the one that counts deferred constraint violations). + // If P1 is zero, the jump is taken if the statement constraint-counter is zero (immediate foreign key constraint violations). + FkIfZero { + if_zero: bool, + target_pc: BranchOffset, + }, } const fn get_insn_virtual_table() -> [InsnFunction; InsnVariants::COUNT] { @@ -1335,6 +1349,8 @@ impl InsnVariants { InsnVariants::MemMax => execute::op_mem_max, InsnVariants::Sequence => execute::op_sequence, InsnVariants::SequenceTest => execute::op_sequence_test, + InsnVariants::FkCounter => execute::op_fk_counter, + InsnVariants::FkIfZero => execute::op_fk_if_zero, } } } diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs index 16695bd0f..4c558a2cc 100644 --- a/core/vdbe/mod.rs +++ b/core/vdbe/mod.rs @@ -313,6 +313,7 @@ pub struct ProgramState { /// This is used when statement in auto-commit mode reseted after previous uncomplete execution - in which case we may need to rollback transaction started on previous attempt /// Note, that MVCC transactions are always explicit - so they do not update auto_txn_cleanup marker pub(crate) auto_txn_cleanup: TxnCleanup, + fk_constraint_counter: isize, } impl ProgramState { @@ -359,6 +360,7 @@ impl ProgramState { op_checkpoint_state: OpCheckpointState::StartCheckpoint, view_delta_state: ViewDeltaCommitState::NotStarted, auto_txn_cleanup: TxnCleanup::None, + fk_constraint_counter: 0, } }