Preserve foreign key clauses when adding columns via ALTER TABLE

Author: ddwalias

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,
                     });
                 },
             )?

core/vdbe/execute.rs

@@ -8561,7 +8561,14 @@ pub fn op_add_column(
     pager: &Arc<Pager>,
     mv_store: Option<&Arc<MvStore>>,
 ) -> Result<InsnFunctionStepResult> {
-    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;

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",

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<ForeignKey>,
     },
     AlterColumn {
         table: String,

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_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);
+} {FOREIGN KEY constraint failed}
+
 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;
-} {{}}
+} {{}}