Skip to content

Commit

Permalink
refactoring Schema to expose functions for database updates (#1256)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
mohs8421 and billy1624 authored Dec 19, 2022
1 parent 17ed715 commit 1f27837
Show file tree
Hide file tree
Showing 2 changed files with 185 additions and 82 deletions.
99 changes: 97 additions & 2 deletions src/entity/relation.rs
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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::{
Expand Down
168 changes: 88 additions & 80 deletions src/schema/entity.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
use crate::{
unpack_table_ref, ActiveEnum, ColumnTrait, ColumnType, DbBackend, EntityTrait, Identity,
Iterable, PrimaryKeyToColumn, PrimaryKeyTrait, RelationTrait, Schema,
ActiveEnum, ColumnTrait, ColumnType, DbBackend, EntityTrait, Iterable, PrimaryKeyToColumn,
PrimaryKeyTrait, RelationTrait, Schema,
};
use sea_query::{
extension::postgres::{Type, TypeCreateStatement},
ColumnDef, ForeignKeyCreateStatement, Iden, Index, IndexCreateStatement, TableCreateStatement,
ColumnDef, Iden, Index, IndexCreateStatement, TableCreateStatement,
};

impl Schema {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -117,39 +162,7 @@ where
let mut stmt = TableCreateStatement::new();

for column in E::Column::iter() {
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();
}
}
}
let mut column_def = column_def_from_entity_column::<E>(column, backend);
stmt.col(&mut column_def);
}

Expand All @@ -166,55 +179,50 @@ where
if relation.is_owner {
continue;
}
let mut foreign_key_stmt = ForeignKeyCreateStatement::new();
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);
stmt.foreign_key(&mut relation.into());
}

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("', '")))
}
Identity::Binary(o1, o2) => {
foreign_key_stmt.to_col(o1);
foreign_key_stmt.to_col(o2);
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();
}
Identity::Ternary(o1, o2, o3) => {
foreign_key_stmt.to_col(o1);
foreign_key_stmt.to_col(o2);
foreign_key_stmt.to_col(o3);
if E::PrimaryKey::iter().count() == 1 {
column_def.primary_key();
}
}
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()
column_def
}

#[cfg(test)]
Expand Down

0 comments on commit 1f27837

Please sign in to comment.