From f504ad1eef41d892925d5924a61bd911e1e567e9 Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Thu, 11 Nov 2021 15:11:52 +0800
Subject: [PATCH] Refactor `Schema`

---
 src/schema/entity.rs      | 20 +++++++++++---------
 src/schema/mod.rs         | 16 ++++++++++++++--
 tests/common/setup/mod.rs |  7 +++++--
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/src/schema/entity.rs b/src/schema/entity.rs
index c24a1385d..8b0ad1444 100644
--- a/src/schema/entity.rs
+++ b/src/schema/entity.rs
@@ -9,27 +9,27 @@ use sea_query::{
 
 impl Schema {
     /// Creates Postgres enums from an Entity. See [TypeCreateStatement] for more details
-    pub fn create_enum_from_entity<E>(entity: E, db_backend: DbBackend) -> Vec<TypeCreateStatement>
+    pub fn create_enum_from_entity<E>(&self, entity: E) -> Vec<TypeCreateStatement>
     where
         E: EntityTrait,
     {
-        create_enum_from_entity(entity, db_backend)
+        create_enum_from_entity(entity, self.backend)
     }
 
     /// Creates a table from an Entity. See [TableCreateStatement] for more details
-    pub fn create_table_from_entity<E>(entity: E, db_backend: DbBackend) -> TableCreateStatement
+    pub fn create_table_from_entity<E>(&self, entity: E) -> TableCreateStatement
     where
         E: EntityTrait,
     {
-        create_table_from_entity(entity, db_backend)
+        create_table_from_entity(entity, self.backend)
     }
 }
 
-pub(crate) fn create_enum_from_entity<E>(_: E, db_backend: DbBackend) -> Vec<TypeCreateStatement>
+pub(crate) fn create_enum_from_entity<E>(_: E, backend: DbBackend) -> Vec<TypeCreateStatement>
 where
     E: EntityTrait,
 {
-    if matches!(db_backend, DbBackend::MySql | DbBackend::Sqlite) {
+    if matches!(backend, DbBackend::MySql | DbBackend::Sqlite) {
         return Vec::new();
     }
     let mut vec = Vec::new();
@@ -52,7 +52,7 @@ where
     vec
 }
 
-pub(crate) fn create_table_from_entity<E>(entity: E, db_backend: DbBackend) -> TableCreateStatement
+pub(crate) fn create_table_from_entity<E>(entity: E, backend: DbBackend) -> TableCreateStatement
 where
     E: EntityTrait,
 {
@@ -61,7 +61,7 @@ where
     for column in E::Column::iter() {
         let orm_column_def = column.def();
         let types = match orm_column_def.col_type {
-            ColumnType::Enum(s, variants) => match db_backend {
+            ColumnType::Enum(s, variants) => match backend {
                 DbBackend::MySql => {
                     ColumnType::Custom(format!("ENUM('{}')", variants.join("', '")))
                 }
@@ -175,8 +175,10 @@ mod tests {
 
     #[test]
     fn test_create_table_from_entity() {
+        let schema = Schema::new(DbBackend::MySql);
         assert_eq!(
-            Schema::create_table_from_entity(CakeFillingPrice, DbBackend::MySql)
+            schema
+                .create_table_from_entity(CakeFillingPrice)
                 .to_string(MysqlQueryBuilder),
             Table::create()
                 .table(CakeFillingPrice)
diff --git a/src/schema/mod.rs b/src/schema/mod.rs
index 3a5ce7345..e89d9cc5e 100644
--- a/src/schema/mod.rs
+++ b/src/schema/mod.rs
@@ -1,5 +1,17 @@
+use crate::DbBackend;
+
 mod entity;
 
-/// This structure defines a schema for a table
+/// This is a helper struct to convert [`EntityTrait`](crate::EntityTrait)
+/// into different [`sea_query`](crate::sea_query) statements.
 #[derive(Debug)]
-pub struct Schema;
+pub struct Schema {
+    backend: DbBackend,
+}
+
+impl Schema {
+    /// Create a helper for a specific database backend
+    pub fn new(backend: DbBackend) -> Self {
+        Self { backend }
+    }
+}
diff --git a/tests/common/setup/mod.rs b/tests/common/setup/mod.rs
index fae61984b..62e06b654 100644
--- a/tests/common/setup/mod.rs
+++ b/tests/common/setup/mod.rs
@@ -108,7 +108,9 @@ where
     }
 
     let expect_stmts: Vec<Statement> = creates.iter().map(|stmt| builder.build(stmt)).collect();
-    let create_from_entity_stmts: Vec<Statement> = Schema::create_enum_from_entity(entity, builder)
+    let schema = Schema::new(builder);
+    let create_from_entity_stmts: Vec<Statement> = schema
+        .create_enum_from_entity(entity)
         .iter()
         .map(|stmt| builder.build(stmt))
         .collect();
@@ -131,8 +133,9 @@ where
     E: EntityTrait,
 {
     let builder = db.get_database_backend();
+    let schema = Schema::new(builder);
     assert_eq!(
-        builder.build(&Schema::create_table_from_entity(entity, builder)),
+        builder.build(&schema.create_table_from_entity(entity)),
         builder.build(create)
     );