diff --git a/core/translate/alter.rs b/core/translate/alter.rs index 4a4d6664e0..da36ca136e 100644 --- a/core/translate/alter.rs +++ b/core/translate/alter.rs @@ -6,7 +6,7 @@ use turso_parser::{ use crate::{ function::{AlterTableFunc, Func}, - schema::{Column, Table, RESERVED_TABLE_PREFIXES}, + schema::{Column, ForeignKey, Table, RESERVED_TABLE_PREFIXES}, translate::{ emitter::Resolver, expr::{walk_expr, WalkControl}, @@ -308,8 +308,63 @@ pub fn translate_alter_table( )); } + let mut column_foreign_keys = Vec::new(); + for constraint in &col_def.constraints { + if let ast::ColumnConstraint::ForeignKey { + clause, + defer_clause, + } = &constraint.constraint + { + let foreign_key = ForeignKey { + child_columns: vec![new_column_name.clone()], + parent_table: normalize_ident(clause.tbl_name.as_str()), + parent_columns: clause + .columns + .iter() + .map(|c| normalize_ident(c.col_name.as_str())) + .collect(), + on_delete: clause + .args + .iter() + .find_map(|arg| { + if let ast::RefArg::OnDelete(act) = arg { + Some(*act) + } else { + None + } + }) + .unwrap_or(ast::RefAct::NoAction), + on_update: clause + .args + .iter() + .find_map(|arg| { + if let ast::RefArg::OnUpdate(act) = arg { + Some(*act) + } else { + None + } + }) + .unwrap_or(ast::RefAct::NoAction), + deferred: match defer_clause { + Some(defer_clause) => { + defer_clause.deferrable + && matches!( + defer_clause.init_deferred, + Some(ast::InitDeferredPred::InitiallyDeferred) + ) + } + None => false, + }, + }; + column_foreign_keys.push(foreign_key); + } + } + // TODO: All quoted ids will be quoted with `[]`, we should store some info from the parsed AST btree.columns.push(column.clone()); + for foreign_key in &column_foreign_keys { + btree.foreign_keys.push(Arc::new(foreign_key.clone())); + } let sql = btree.to_sql(); let mut escaped = String::with_capacity(sql.len()); @@ -351,6 +406,7 @@ pub fn translate_alter_table( program.emit_insn(Insn::AddColumn { table: table_name.to_owned(), column, + foreign_keys: column_foreign_keys, }); }, )? diff --git a/core/vdbe/execute.rs b/core/vdbe/execute.rs index 1ddb215535..523c360f50 100644 --- a/core/vdbe/execute.rs +++ b/core/vdbe/execute.rs @@ -8561,7 +8561,14 @@ pub fn op_add_column( pager: &Arc, mv_store: Option<&Arc>, ) -> Result { - load_insn!(AddColumn { table, column }, insn); + load_insn!( + AddColumn { + table, + column, + foreign_keys + }, + insn + ); let conn = program.connection.clone(); @@ -8578,7 +8585,10 @@ pub fn op_add_column( }; let btree = Arc::make_mut(btree); - btree.columns.push(column.clone()) + btree.columns.push(column.clone()); + for fk in foreign_keys { + btree.foreign_keys.push(Arc::new(fk.clone())); + } }); state.pc += 1; diff --git a/core/vdbe/explain.rs b/core/vdbe/explain.rs index ccaca542db..e870b6651d 100644 --- a/core/vdbe/explain.rs +++ b/core/vdbe/explain.rs @@ -1742,14 +1742,21 @@ pub fn insn_to_row( 0, format!("drop_column({table}, {column_index})"), ), - Insn::AddColumn { table, column } => ( + Insn::AddColumn { + table, + column, + foreign_keys, + } => ( "AddColumn", 0, 0, 0, Value::build_text(""), 0, - format!("add_column({table}, {column:?})"), + format!( + "add_column({table}, {column:?}, fks={})", + foreign_keys.len() + ), ), Insn::AlterColumn { table, column_index, definition: column, rename } => ( "AlterColumn", diff --git a/core/vdbe/insn.rs b/core/vdbe/insn.rs index 98693d7187..84cce06332 100644 --- a/core/vdbe/insn.rs +++ b/core/vdbe/insn.rs @@ -5,7 +5,7 @@ use std::{ use super::{execute, AggFunc, BranchOffset, CursorID, FuncCtx, InsnFunction, PageIdx}; use crate::{ - schema::{Affinity, BTreeTable, Column, Index}, + schema::{Affinity, BTreeTable, Column, ForeignKey, Index}, storage::{pager::CreateBTreeFlags, wal::CheckpointMode}, translate::{collate::CollationSeq, emitter::TransactionMode}, types::KeyInfo, @@ -1149,6 +1149,7 @@ pub enum Insn { AddColumn { table: String, column: Column, + foreign_keys: Vec, }, AlterColumn { table: String, diff --git a/testing/alter_table.test b/testing/alter_table.test index e858d7e7b0..9ea3e3d9cf 100755 --- a/testing/alter_table.test +++ b/testing/alter_table.test @@ -90,6 +90,27 @@ do_execsql_test_on_specific_db {:memory:} alter-table-add-quoted-column { "CREATE TABLE test (a, [b c])" } +do_execsql_test_on_specific_db {:memory:} alter-table-add-column-foreign-key { + CREATE TABLE parent(a PRIMARY KEY); + CREATE TABLE child(a); + + ALTER TABLE child ADD COLUMN b REFERENCES parent(a); + SELECT sql FROM sqlite_schema WHERE name = 'child'; + PRAGMA foreign_keys = ON; + INSERT INTO parent VALUES (1); + INSERT INTO child VALUES (1, 1); +} { + "CREATE TABLE child (a, b, FOREIGN KEY (b) REFERENCES parent(a))" +} + +do_execsql_test_in_memory_any_error alter-table-add-column-foreign-key-enforced { + PRAGMA foreign_keys = ON; + CREATE TABLE parent(a PRIMARY KEY); + CREATE TABLE child(a); + ALTER TABLE child ADD COLUMN b REFERENCES parent(a); + INSERT INTO child VALUES (2, 2); +} + do_execsql_test_on_specific_db {:memory:} alter-table-drop-column { CREATE TABLE t (a, b); INSERT INTO t VALUES (1, 1), (2, 2), (3, 3); @@ -219,4 +240,4 @@ do_execsql_test_on_specific_db {:memory:} drop-column-regression { ALTER TABLE t ADD COLUMN col3 BLOB; ALTER TABLE t DROP COLUMN col1; SELECT col3 FROM t; -} {{}} \ No newline at end of file +} {{}}