refactoring Schema to expose functions for database updates (#1256)

* extracting get_column_def from create_table_from_entity to make it available for database upgrade processes.

* Align code example formatting

* Converting the foreign key related code from create_table_from_entity into From<RelationDef> implementations to make its usage easier in different context, like updating a database.

* Refactor

* Fixup

Co-authored-by: Billy Chan <ccw.billy.123@gmail.com>

src/entity/relation.rs

@@ -1,6 +1,9 @@
-use crate::{EntityTrait, Identity, IdentityOf, Iterable, QuerySelect, Select};
+use crate::{unpack_table_ref, EntityTrait, Identity, IdentityOf, Iterable, QuerySelect, Select};
 use core::marker::PhantomData;
-use sea_query::{Alias, Condition, DynIden, JoinType, SeaRc, TableRef};
+use sea_query::{
+    Alias, Condition, DynIden, ForeignKeyCreateStatement, JoinType, SeaRc, TableForeignKey,
+    TableRef,
+};
 use std::fmt::Debug;
 
 /// Defines the type of relationship
@@ -309,6 +312,98 @@ where
     }
 }
 
+macro_rules! set_foreign_key_stmt {
+    ( $relation: ident, $foreign_key: ident ) => {
+        let from_cols: Vec<String> = match $relation.from_col {
+            Identity::Unary(o1) => vec![o1],
+            Identity::Binary(o1, o2) => vec![o1, o2],
+            Identity::Ternary(o1, o2, o3) => vec![o1, o2, o3],
+        }
+        .into_iter()
+        .map(|col| {
+            let col_name = col.to_string();
+            $foreign_key.from_col(col);
+            col_name
+        })
+        .collect();
+        match $relation.to_col {
+            Identity::Unary(o1) => {
+                $foreign_key.to_col(o1);
+            }
+            Identity::Binary(o1, o2) => {
+                $foreign_key.to_col(o1);
+                $foreign_key.to_col(o2);
+            }
+            Identity::Ternary(o1, o2, o3) => {
+                $foreign_key.to_col(o1);
+                $foreign_key.to_col(o2);
+                $foreign_key.to_col(o3);
+            }
+        }
+        if let Some(action) = $relation.on_delete {
+            $foreign_key.on_delete(action);
+        }
+        if let Some(action) = $relation.on_update {
+            $foreign_key.on_update(action);
+        }
+        let name = if let Some(name) = $relation.fk_name {
+            name
+        } else {
+            let from_tbl = unpack_table_ref(&$relation.from_tbl);
+            format!("fk-{}-{}", from_tbl.to_string(), from_cols.join("-"))
+        };
+        $foreign_key.name(&name);
+    };
+}
+
+impl From<RelationDef> for ForeignKeyCreateStatement {
+    fn from(relation: RelationDef) -> Self {
+        let mut foreign_key_stmt = Self::new();
+        set_foreign_key_stmt!(relation, foreign_key_stmt);
+        foreign_key_stmt
+            .from_tbl(unpack_table_ref(&relation.from_tbl))
+            .to_tbl(unpack_table_ref(&relation.to_tbl))
+            .take()
+    }
+}
+
+/// Creates a column definition for example to update a table.
+/// ```
+/// use sea_query::{Alias, IntoIden, MysqlQueryBuilder, TableAlterStatement, TableRef};
+/// use sea_orm::{EnumIter, Iden, Identity, PrimaryKeyTrait, RelationDef, RelationTrait, RelationType};
+///
+/// let relation = RelationDef {
+///     rel_type: RelationType::HasOne,
+///     from_tbl: TableRef::Table(Alias::new("foo").into_iden()),
+///     to_tbl: TableRef::Table(Alias::new("bar").into_iden()),
+///     from_col: Identity::Unary(Alias::new("bar_id").into_iden()),
+///     to_col: Identity::Unary(Alias::new("bar_id").into_iden()),
+///     is_owner: false,
+///     on_delete: None,
+///     on_update: None,
+///     on_condition: None,
+///     fk_name: Some("foo-bar".to_string()),
+/// };
+///
+/// let mut alter_table = TableAlterStatement::new()
+///     .table(TableRef::Table(Alias::new("foo").into_iden()))
+///     .add_foreign_key(&mut relation.into()).take();
+/// assert_eq!(
+///     alter_table.to_string(MysqlQueryBuilder::default()),
+///     "ALTER TABLE `foo` ADD CONSTRAINT `foo-bar` FOREIGN KEY (`bar_id`) REFERENCES `bar` (`bar_id`)"
+/// );
+/// ```
+impl From<RelationDef> for TableForeignKey {
+    fn from(relation: RelationDef) -> Self {
+        let mut foreign_key = Self::new();
+        set_foreign_key_stmt!(relation, foreign_key);
+        foreign_key
+            .from_tbl(unpack_table_ref(&relation.from_tbl))
+            .to_tbl(unpack_table_ref(&relation.to_tbl))
+            .take()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{

src/schema/entity.rs

@@ -1,10 +1,10 @@
 use crate::{
-    unpack_table_ref, ActiveEnum, ColumnTrait, ColumnType, DbBackend, EntityTrait, Identity,
+    ActiveEnum, ColumnTrait, ColumnType, DbBackend, EntityTrait, Iterable, PrimaryKeyToColumn,
-    Iterable, PrimaryKeyToColumn, PrimaryKeyTrait, RelationTrait, Schema,
+    PrimaryKeyTrait, RelationTrait, Schema,
 };
 use sea_query::{
     extension::postgres::{Type, TypeCreateStatement},
-    ColumnDef, ForeignKeyCreateStatement, Iden, Index, IndexCreateStatement, TableCreateStatement,
+    ColumnDef, Iden, Index, IndexCreateStatement, TableCreateStatement,
 };
 
 impl Schema {
@@ -40,6 +40,51 @@ impl Schema {
     {
         create_index_from_entity(entity, self.backend)
     }
+
+    /// Creates a column definition for example to update a table.
+    /// ```
+    /// use crate::sea_orm::IdenStatic;
+    /// use sea_orm::{
+    ///     ActiveModelBehavior, ColumnDef, ColumnTrait, ColumnType, DbBackend, EntityName,
+    ///     EntityTrait, EnumIter, PrimaryKeyTrait, RelationDef, RelationTrait, Schema,
+    /// };
+    /// use sea_orm_macros::{DeriveEntityModel, DerivePrimaryKey};
+    /// use sea_query::{MysqlQueryBuilder, TableAlterStatement};
+    ///
+    /// #[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+    /// #[sea_orm(table_name = "posts")]
+    /// pub struct Model {
+    ///     #[sea_orm(primary_key)]
+    ///     pub id: u32,
+    ///     pub title: String,
+    /// }
+    ///
+    /// #[derive(Copy, Clone, Debug, EnumIter)]
+    /// pub enum Relation {}
+    /// impl RelationTrait for Relation {
+    ///     fn def(&self) -> RelationDef {
+    ///         panic!("No RelationDef")
+    ///     }
+    /// }
+    /// impl ActiveModelBehavior for ActiveModel {}
+    ///
+    /// let schema = Schema::new(DbBackend::MySql);
+    ///
+    /// let mut alter_table = TableAlterStatement::new()
+    ///     .table(Entity)
+    ///     .add_column(&mut schema.get_column_def::<Entity>(Column::Title))
+    ///     .take();
+    /// assert_eq!(
+    ///     alter_table.to_string(MysqlQueryBuilder::default()),
+    ///     "ALTER TABLE `posts` ADD COLUMN `title` varchar(255) NOT NULL"
+    /// );
+    /// ```
+    pub fn get_column_def<E>(&self, column: E::Column) -> ColumnDef
+    where
+        E: EntityTrait,
+    {
+        column_def_from_entity_column::<E>(column, self.backend)
+    }
 }
 
 pub(crate) fn create_enum_from_active_enum<A>(backend: DbBackend) -> TypeCreateStatement
@@ -117,39 +162,7 @@ where
     let mut stmt = TableCreateStatement::new();
 
     for column in E::Column::iter() {
-        let orm_column_def = column.def();
+        let mut column_def = column_def_from_entity_column::<E>(column, backend);
-        let types = match orm_column_def.col_type {
-            ColumnType::Enum { name, variants } => match backend {
-                DbBackend::MySql => {
-                    let variants: Vec<String> = variants.iter().map(|v| v.to_string()).collect();
-                    ColumnType::Custom(format!("ENUM('{}')", variants.join("', '")))
-                }
-                DbBackend::Postgres => ColumnType::Custom(name.to_string()),
-                DbBackend::Sqlite => ColumnType::Text,
-            }
-            .into(),
-            _ => orm_column_def.col_type.into(),
-        };
-        let mut column_def = ColumnDef::new_with_type(column, types);
-        if !orm_column_def.null {
-            column_def.not_null();
-        }
-        if orm_column_def.unique {
-            column_def.unique_key();
-        }
-        if let Some(value) = orm_column_def.default_value {
-            column_def.default(value);
-        }
-        for primary_key in E::PrimaryKey::iter() {
-            if column.to_string() == primary_key.into_column().to_string() {
-                if E::PrimaryKey::auto_increment() {
-                    column_def.auto_increment();
-                }
-                if E::PrimaryKey::iter().count() == 1 {
-                    column_def.primary_key();
-                }
-            }
-        }
         stmt.col(&mut column_def);
     }
 
@@ -166,57 +179,52 @@ where
         if relation.is_owner {
             continue;
         }
-        let mut foreign_key_stmt = ForeignKeyCreateStatement::new();
+        stmt.foreign_key(&mut relation.into());
-        let from_tbl = unpack_table_ref(&relation.from_tbl);
-        let to_tbl = unpack_table_ref(&relation.to_tbl);
-        let from_cols: Vec<String> = match relation.from_col {
-            Identity::Unary(o1) => vec![o1],
-            Identity::Binary(o1, o2) => vec![o1, o2],
-            Identity::Ternary(o1, o2, o3) => vec![o1, o2, o3],
-        }
-        .into_iter()
-        .map(|col| {
-            let col_name = col.to_string();
-            foreign_key_stmt.from_col(col);
-            col_name
-        })
-        .collect();
-        match relation.to_col {
-            Identity::Unary(o1) => {
-                foreign_key_stmt.to_col(o1);
-            }
-            Identity::Binary(o1, o2) => {
-                foreign_key_stmt.to_col(o1);
-                foreign_key_stmt.to_col(o2);
-            }
-            Identity::Ternary(o1, o2, o3) => {
-                foreign_key_stmt.to_col(o1);
-                foreign_key_stmt.to_col(o2);
-                foreign_key_stmt.to_col(o3);
-            }
-        }
-        if let Some(action) = relation.on_delete {
-            foreign_key_stmt.on_delete(action);
-        }
-        if let Some(action) = relation.on_update {
-            foreign_key_stmt.on_update(action);
-        }
-        let name = if let Some(name) = relation.fk_name {
-            name
-        } else {
-            format!("fk-{}-{}", from_tbl.to_string(), from_cols.join("-"))
-        };
-        stmt.foreign_key(
-            foreign_key_stmt
-                .name(&name)
-                .from_tbl(from_tbl)
-                .to_tbl(to_tbl),
-        );
     }
 
     stmt.table(entity.table_ref()).take()
 }
 
+fn column_def_from_entity_column<E>(column: E::Column, backend: DbBackend) -> ColumnDef
+where
+    E: EntityTrait,
+{
+    let orm_column_def = column.def();
+    let types = match orm_column_def.col_type {
+        ColumnType::Enum { name, variants } => match backend {
+            DbBackend::MySql => {
+                let variants: Vec<String> = variants.iter().map(|v| v.to_string()).collect();
+                ColumnType::Custom(format!("ENUM('{}')", variants.join("', '")))
+            }
+            DbBackend::Postgres => ColumnType::Custom(name.to_string()),
+            DbBackend::Sqlite => ColumnType::Text,
+        }
+        .into(),
+        _ => orm_column_def.col_type.into(),
+    };
+    let mut column_def = ColumnDef::new_with_type(column, types);
+    if !orm_column_def.null {
+        column_def.not_null();
+    }
+    if orm_column_def.unique {
+        column_def.unique_key();
+    }
+    if let Some(value) = orm_column_def.default_value {
+        column_def.default(value);
+    }
+    for primary_key in E::PrimaryKey::iter() {
+        if column.to_string() == primary_key.into_column().to_string() {
+            if E::PrimaryKey::auto_increment() {
+                column_def.auto_increment();
+            }
+            if E::PrimaryKey::iter().count() == 1 {
+                column_def.primary_key();
+            }
+        }
+    }
+    column_def
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{sea_query::*, tests_cfg::*, DbBackend, EntityName, Schema};