Add foreign keys option to create table migration

README.md

@@ -107,7 +107,7 @@ Every action has a `type`. The supported types are detailed below.
 
 ### Create table
 
-The `create_table` action will create a new table with the specified columns and indices.
+The `create_table` action will create a new table with the specified columns, indices and constraints.
 
 *Example: creating a `customers` table with a few columns and a primary key*
 
@@ -132,6 +132,37 @@ primary_key = "id"
 	default = "'PLACEHOLDER'" 
 ```
 
+*Example: creating `users` and `items` tables with a foreign key between them*
+
+```toml
+[[actions]]
+type = "create_table"
+table = "users"
+primary_key = "id"
+
+	[[actions.columns]]
+	name = "id"
+	type = "SERIAL"
+
+[[actions]]
+type = "create_table"
+table = "items"
+primary_key = "id"
+
+	[[actions.columns]]
+	name = "id"
+	type = "SERIAL"
+
+	[[actions.columns]]
+	name = "user_id"
+	type = "INTEGER"
+
+	[[actions.foreign_keys]]
+	columns = ["user_id"]
+	referenced_table = "users"
+	referenced_columns = ["id"]
+```
+
 ### Add column
 
 The `add_column` action will add a new column to an existing table. You can optionally provide an `up` setting. This should be an SQL expression which will be run for all existing rows to backfill the new column.

src/migrations/create_table.rs

@@ -10,6 +10,14 @@ pub struct CreateTable {
     pub name: String,
     pub columns: Vec<Column>,
     pub primary_key: Option<String>,
+    pub foreign_keys: Vec<ForeignKey>,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct ForeignKey {
+    pub columns: Vec<String>,
+    pub referenced_table: String,
+    pub referenced_columns: Vec<String>,
 }
 
 #[typetag::serde(name = "create_table")]
@@ -42,14 +50,22 @@ impl Action for CreateTable {
             definition_rows.push(format!("PRIMARY KEY ({})", column));
         }
 
-        let query = format!(
+        for foreign_key in &self.foreign_keys {
+            definition_rows.push(format!(
+                "FOREIGN KEY ({columns}) REFERENCES {table} ({referenced_columns})",
+                columns = foreign_key.columns.join(", "),
+                table = foreign_key.referenced_table,
+                referenced_columns = foreign_key.referenced_columns.join(", "),
+            ));
+        }
+
+        db.run(&format!(
             "CREATE TABLE {} (
                 {}
             )",
             self.name,
             definition_rows.join(",\n"),
-        );
+        ))?;
-        db.run(&query)?;
         Ok(())
     }
 

src/migrations/mod.rs

@@ -7,7 +7,7 @@ mod common;
 pub use common::Column;
 
 mod create_table;
-pub use create_table::CreateTable;
+pub use create_table::{CreateTable, ForeignKey};
 
 mod alter_column;
 pub use alter_column::{AlterColumn, ColumnChanges};

tests/add_column.rs

@@ -10,6 +10,7 @@ fn add_column() {
     let create_users_table = Migration::new("create_users_table", None).with_action(CreateTable {
         name: "users".to_string(),
         primary_key: None,
+        foreign_keys: vec![],
         columns: vec![
             Column {
                 name: "id".to_string(),
@@ -124,6 +125,7 @@ fn add_column_nullable() {
     let create_users_table = Migration::new("create_users_table", None).with_action(CreateTable {
         name: "users".to_string(),
         primary_key: None,
+        foreign_keys: vec![],
         columns: vec![Column {
             name: "id".to_string(),
             data_type: "SERIAL".to_string(),
@@ -217,13 +219,14 @@ fn add_column_with_default() {
 
     let create_users_table = Migration::new("create_users_table", None).with_action(CreateTable {
         name: "users".to_string(),
+        primary_key: None,
+        foreign_keys: vec![],
         columns: vec![Column {
             name: "id".to_string(),
             data_type: "SERIAL".to_string(),
             nullable: true,
             default: None,
         }],
-        primary_key: None,
     });
     let add_name_column =
         Migration::new("add_name_column_with_default", None).with_action(AddColumn {

tests/alter_column.rs

@@ -10,6 +10,7 @@ fn alter_column_data() {
     let create_users_table = Migration::new("create_users_table", None).with_action(CreateTable {
         name: "users".to_string(),
         primary_key: None,
+        foreign_keys: vec![],
         columns: vec![
             Column {
                 name: "id".to_string(),
@@ -111,6 +112,7 @@ fn alter_column_set_not_null() {
     let create_users_table = Migration::new("create_users_table", None).with_action(CreateTable {
         name: "users".to_string(),
         primary_key: None,
+        foreign_keys: vec![],
         columns: vec![
             Column {
                 name: "id".to_string(),
@@ -212,6 +214,7 @@ fn alter_column_rename() {
     let create_users_table = Migration::new("create_users_table", None).with_action(CreateTable {
         name: "users".to_string(),
         primary_key: None,
+        foreign_keys: vec![],
         columns: vec![
             Column {
                 name: "id".to_string(),

tests/create_table.rs

@@ -1,5 +1,5 @@
 use reshape::{
-    migrations::{Column, CreateTable, Migration},
+    migrations::{Column, CreateTable, ForeignKey, Migration},
     Status,
 };
 
@@ -13,6 +13,7 @@ fn create_table() {
         Migration::new("create_users_table", None).with_action(CreateTable {
             name: "users".to_string(),
             primary_key: Some("id".to_string()),
+            foreign_keys: vec![],
             columns: vec![
                 Column {
                     name: "id".to_string(),
@@ -117,3 +118,80 @@ fn create_table() {
 
     assert_eq!(vec!["id"], primary_key_columns);
 }
+
+#[test]
+fn create_table_with_foreign_keys() {
+    let (mut reshape, mut db, _) = common::setup();
+
+    let create_table_migration =
+        Migration::new("create_users_table", None).with_action(CreateTable {
+            name: "users".to_string(),
+            primary_key: Some("id".to_string()),
+            foreign_keys: vec![],
+            columns: vec![Column {
+                name: "id".to_string(),
+                data_type: "SERIAL".to_string(),
+                nullable: true, // Will be ignored by Postgres as the column is a SERIAL
+                default: None,
+            }],
+        });
+
+    let create_second_table_migration =
+        Migration::new("create_items_table", None).with_action(CreateTable {
+            name: "items".to_string(),
+            primary_key: None,
+            foreign_keys: vec![ForeignKey {
+                columns: vec!["user_id".to_string()],
+                referenced_table: "users".to_string(),
+                referenced_columns: vec!["id".to_string()],
+            }],
+            columns: vec![Column {
+                name: "user_id".to_string(),
+                data_type: "INTEGER".to_string(),
+                nullable: false,
+                default: None,
+            }],
+        });
+
+    reshape
+        .migrate(vec![
+            create_table_migration.clone(),
+            create_second_table_migration.clone(),
+        ])
+        .unwrap();
+
+    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';
+        ",
+            &[],
+        )
+        .unwrap()
+        .iter()
+        .map(|row| {
+            (
+                row.get("column_name"),
+                row.get("foreign_table_name"),
+                row.get("foreign_column_name"),
+            )
+        })
+        .collect();
+
+    assert_eq!(
+        vec![("user_id".to_string(), "users".to_string(), "id".to_string())],
+        foreign_key_columns
+    );
+}

tests/remove_column.rs

@@ -10,6 +10,7 @@ fn remove_column() {
         Migration::new("create_users_table", None).with_action(CreateTable {
             name: "users".to_string(),
             primary_key: None,
+            foreign_keys: vec![],
             columns: vec![
                 Column {
                     name: "id".to_string(),