Add new add_foreign_key action

README.md

@@ -21,6 +21,7 @@ Designed for Postgres 12 and later.
 		- [Create table](#create-table)
 		- [Rename table](#rename-table)
 		- [Remove table](#remove-table)
+		- [Add foreign key](#add-foreign-key)
 	- [Columns](#columns)
 		- [Add column](#add-column)
 		- [Alter column](#alter-column)
@@ -243,6 +244,23 @@ type = "remove_table"
 table = "users"
 ```
 
+#### Add foreign key
+
+The `add_foreign_key` action will add a foreign key between two existing tables. The migration will fail if the existing column values aren't valid references.
+
+*Example: create foreign key from `items` to `users` table*
+
+```toml
+[[actions]]
+type = "add_foreign_key"
+table = "items"
+
+	[actions.foreign_key]
+	columns = ["user_id"]
+	referenced_table = "users"
+	referenced_columns = ["id"]
+```
+
 ### Columns
 
 #### Add column

src/migrations/add_foreign_key.rs

@@ -0,0 +1,122 @@
+use super::{common::ForeignKey, Action, MigrationContext};
+use crate::{
+    db::{Conn, Transaction},
+    schema::Schema,
+};
+use anyhow::Context;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct AddForeignKey {
+    pub table: String,
+    foreign_key: ForeignKey,
+}
+
+#[typetag::serde(name = "add_foreign_key")]
+impl Action for AddForeignKey {
+    fn describe(&self) -> String {
+        format!(
+            "Adding foreign key from table \"{}\" to \"{}\"",
+            self.table, self.foreign_key.referenced_table
+        )
+    }
+
+    fn run(
+        &self,
+        ctx: &MigrationContext,
+        db: &mut dyn Conn,
+        schema: &Schema,
+    ) -> anyhow::Result<()> {
+        let table = schema.get_table(db, &self.table)?;
+        let referenced_table = schema.get_table(db, &self.foreign_key.referenced_table)?;
+
+        // Add quotes around all column names
+        let columns: Vec<String> = table
+            .real_column_names(&self.foreign_key.columns)
+            .map(|col| format!("\"{}\"", col))
+            .collect();
+        let referenced_columns: Vec<String> = referenced_table
+            .real_column_names(&self.foreign_key.referenced_columns)
+            .map(|col| format!("\"{}\"", col))
+            .collect();
+
+        // Create foreign key but set is as NOT VALID.
+        // This means the foreign key will be enforced for inserts and updates
+        // but the existing data won't be checked, that would cause a long-lived lock.
+        db.run(&format!(
+            r#"
+            ALTER TABLE "{table}"
+            ADD CONSTRAINT {constraint_name}
+            FOREIGN KEY ({columns})
+            REFERENCES {referenced_table} ({referenced_columns})
+            NOT VALID
+            "#,
+            table = table.real_name,
+            constraint_name = self.temp_constraint_name(ctx),
+            columns = columns.join(", "),
+            referenced_table = referenced_table.real_name,
+            referenced_columns = referenced_columns.join(", "),
+        ))
+        .context("failed to create foreign key")?;
+
+        db.run(&format!(
+            r#"
+            ALTER TABLE "{table}"
+            VALIDATE CONSTRAINT "{constraint_name}"
+            "#,
+            table = table.real_name,
+            constraint_name = self.temp_constraint_name(ctx),
+        ))
+        .context("failed to validate foreign key")?;
+
+        Ok(())
+    }
+
+    fn complete<'a>(
+        &self,
+        ctx: &MigrationContext,
+        db: &'a mut dyn Conn,
+    ) -> anyhow::Result<Option<Transaction<'a>>> {
+        db.run(&format!(
+            r#"
+            ALTER TABLE {table}
+            RENAME CONSTRAINT {temp_constraint_name} TO {constraint_name}
+            "#,
+            table = self.table,
+            temp_constraint_name = self.temp_constraint_name(ctx),
+            constraint_name = self.final_constraint_name(),
+        ))
+        .context("failed to rename temporary constraint")?;
+        Ok(None)
+    }
+
+    fn update_schema(&self, _ctx: &MigrationContext, _schema: &mut Schema) {}
+
+    fn abort(&self, ctx: &MigrationContext, db: &mut dyn Conn) -> anyhow::Result<()> {
+        db.run(&format!(
+            r#"
+            ALTER TABLE "{table}"
+            DROP CONSTRAINT IF EXISTS "{constraint_name}"
+            "#,
+            table = self.table,
+            constraint_name = self.temp_constraint_name(ctx),
+        ))
+        .context("failed to validate foreign key")?;
+
+        Ok(())
+    }
+}
+
+impl AddForeignKey {
+    fn temp_constraint_name(&self, ctx: &MigrationContext) -> String {
+        format!("{}_temp_fkey", ctx.prefix())
+    }
+
+    fn final_constraint_name(&self) -> String {
+        format!(
+            "{table}_{columns}_fkey",
+            table = self.table,
+            columns = self.foreign_key.columns.join("_")
+        )
+    }
+}

src/migrations/common.rs

@@ -19,6 +19,13 @@ fn nullable_default() -> bool {
     true
 }
 
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct ForeignKey {
+    pub columns: Vec<String>,
+    pub referenced_table: String,
+    pub referenced_columns: Vec<String>,
+}
+
 #[derive(Debug)]
 struct PostgresRawValue {
     bytes: Vec<u8>,

src/migrations/mod.rs

@@ -10,7 +10,7 @@ mod common;
 pub use common::Column;
 
 mod create_table;
-pub use create_table::{CreateTable, ForeignKey};
+pub use create_table::CreateTable;
 
 mod alter_column;
 pub use alter_column::{AlterColumn, ColumnChanges};
@@ -39,6 +39,9 @@ pub use create_enum::CreateEnum;
 mod remove_enum;
 pub use remove_enum::RemoveEnum;
 
+mod add_foreign_key;
+pub use add_foreign_key::AddForeignKey;
+
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Migration {
     pub name: String,

src/schema.rs

@@ -1,8 +1,5 @@
 use crate::db::Conn;
-use std::{
-    collections::{HashMap, HashSet},
-    path::Iter,
-};
+use std::collections::{HashMap, HashSet};
 
 // Schema tracks changes made to tables and columns during a migration.
 // These changes are not applied until the migration is completed but

tests/add_foreign_key.rs

@@ -0,0 +1,166 @@
+mod common;
+use common::Test;
+
+#[test]
+fn add_foreign_key() {
+    let mut test = Test::new("Add foreign key");
+
+    test.first_migration(
+        r#"
+        name = "create_user_table"
+
+        [[actions]]
+        type = "create_table"
+        name = "users"
+        primary_key = ["id"]
+
+            [[actions.columns]]
+            name = "id"
+            type = "INTEGER"
+        
+        [[actions]]
+        type = "create_table"
+        name = "items"
+        primary_key = ["id"]
+
+            [[actions.columns]]
+            name = "id"
+            type = "INTEGER"
+
+            [[actions.columns]]
+            name = "user_id"
+            type = "INTEGER"
+        "#,
+    );
+
+    test.second_migration(
+        r#"
+        name = "add_foreign_key"
+
+        [[actions]]
+        type = "add_foreign_key"
+        table = "items"
+
+            [actions.foreign_key]
+            columns = ["user_id"]
+            referenced_table = "users"
+            referenced_columns = ["id"]
+        "#,
+    );
+
+    test.after_first(|db| {
+        // Insert some test users
+        db.simple_query("INSERT INTO users (id) VALUES (1), (2)")
+            .unwrap();
+    });
+
+    test.intermediate(|db, _| {
+        // Ensure items can be inserted if they reference valid users
+        db.simple_query("INSERT INTO items (id, user_id) VALUES (1, 1), (2, 2)")
+            .unwrap();
+
+        // Ensure items can't be inserted if they don't reference valid users
+        let result = db.simple_query("INSERT INTO items (id, user_id) VALUES (3, 3)");
+        assert!(result.is_err(), "expected insert to fail");
+    });
+
+    test.after_completion(|db| {
+        // Ensure items can be inserted if they reference valid users
+        db.simple_query("INSERT INTO items (id, user_id) VALUES (3, 1), (4, 2)")
+            .unwrap();
+
+        // Ensure items can't be inserted if they don't reference valid users
+        let result = db.simple_query("INSERT INTO items (id, user_id) VALUES (5, 3)");
+        assert!(result.is_err(), "expected insert to fail");
+
+        // Ensure foreign key exists with the right name
+        let foreign_key_name: Option<String> = db
+            .query(
+                "
+                SELECT tc.constraint_name
+                FROM information_schema.table_constraints AS tc 
+                WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name='items';
+                ",
+                &[],
+            )
+            .unwrap()
+            .first()
+            .map(|row| row.get(0));
+        assert_eq!(Some("items_user_id_fkey".to_string()), foreign_key_name);
+    });
+
+    test.after_abort(|db| {
+        // Ensure foreign key doesn't exist
+        let fk_does_not_exist = db
+            .query(
+                "
+                SELECT tc.constraint_name
+                FROM information_schema.table_constraints AS tc 
+                WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name='items';
+                ",
+                &[],
+            )
+            .unwrap()
+            .is_empty();
+        assert!(fk_does_not_exist);
+    });
+
+    test.run()
+}
+
+#[test]
+fn add_invalid_foreign_key() {
+    let mut test = Test::new("Add invalid foreign key");
+
+    test.first_migration(
+        r#"
+        name = "create_user_table"
+
+        [[actions]]
+        type = "create_table"
+        name = "users"
+        primary_key = ["id"]
+
+            [[actions.columns]]
+            name = "id"
+            type = "INTEGER"
+        
+        [[actions]]
+        type = "create_table"
+        name = "items"
+        primary_key = ["id"]
+
+            [[actions.columns]]
+            name = "id"
+            type = "INTEGER"
+
+            [[actions.columns]]
+            name = "user_id"
+            type = "INTEGER"
+        "#,
+    );
+
+    test.second_migration(
+        r#"
+        name = "add_foreign_key"
+
+        [[actions]]
+        type = "add_foreign_key"
+        table = "items"
+
+            [actions.foreign_key]
+            columns = ["user_id"]
+            referenced_table = "users"
+            referenced_columns = ["id"]
+        "#,
+    );
+
+    test.after_first(|db| {
+        // Insert some items which don't reference a valid user
+        db.simple_query("INSERT INTO items (id, user_id) VALUES (1, 1), (2, 2)")
+            .unwrap();
+    });
+
+    test.expect_failure();
+    test.run()
+}

tests/create_table.rs

@@ -163,20 +163,20 @@ fn create_table_with_foreign_keys() {
         let foreign_key_columns: Vec<(String, String, String)> = db
             .query(
                 "
-            SELECT
-                kcu.column_name, 
-                ccu.table_name AS foreign_table_name,
-                ccu.column_name AS foreign_column_name 
-            FROM 
-                information_schema.table_constraints AS tc 
-                JOIN information_schema.key_column_usage AS kcu
-                ON tc.constraint_name = kcu.constraint_name
-                AND tc.table_schema = kcu.table_schema
-                JOIN information_schema.constraint_column_usage AS ccu
-                ON ccu.constraint_name = tc.constraint_name
-                AND ccu.table_schema = tc.table_schema
-            WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name='items';
-            ",
+                SELECT
+                    kcu.column_name, 
+                    ccu.table_name AS foreign_table_name,
+                    ccu.column_name AS foreign_column_name 
+                FROM 
+                    information_schema.table_constraints AS tc 
+                    JOIN information_schema.key_column_usage AS kcu
+                    ON tc.constraint_name = kcu.constraint_name
+                    AND tc.table_schema = kcu.table_schema
+                    JOIN information_schema.constraint_column_usage AS ccu
+                    ON ccu.constraint_name = tc.constraint_name
+                    AND ccu.table_schema = tc.table_schema
+                WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name='items';
+                ",
                 &[],
             )
             .unwrap()