From 3e7220aad0c91f3c6de64433e0c923b01139f581 Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Tue, 29 Jun 2021 16:59:40 +0800
Subject: [PATCH] Add SchemaBuilderBackend

---
 src/database/connection.rs | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/database/connection.rs b/src/database/connection.rs
index 95911c29..9727125b 100644
--- a/src/database/connection.rs
+++ b/src/database/connection.rs
@@ -1,6 +1,7 @@
 use crate::{ExecErr, ExecResult, QueryErr, QueryResult, Statement, Transaction};
 use sea_query::{
-    MysqlQueryBuilder, PostgresQueryBuilder, QueryStatementBuilder, SqliteQueryBuilder,
+    MysqlQueryBuilder, PostgresQueryBuilder, QueryStatementBuilder, SchemaStatementBuilder,
+    SqliteQueryBuilder,
 };
 use std::{error::Error, fmt};
 
@@ -22,6 +23,12 @@ pub enum QueryBuilderBackend {
     Sqlite,
 }
 
+pub enum SchemaBuilderBackend {
+    MySql,
+    Postgres,
+    Sqlite,
+}
+
 #[derive(Debug)]
 pub struct ConnectionErr;
 
@@ -70,6 +77,18 @@ impl DatabaseConnection {
         }
     }
 
+    pub fn get_schema_builder_backend(&self) -> SchemaBuilderBackend {
+        match self {
+            #[cfg(feature = "sqlx-mysql")]
+            DatabaseConnection::SqlxMySqlPoolConnection(_) => SchemaBuilderBackend::MySql,
+            #[cfg(feature = "sqlx-sqlite")]
+            DatabaseConnection::SqlxSqlitePoolConnection(_) => SchemaBuilderBackend::Sqlite,
+            #[cfg(feature = "mock")]
+            DatabaseConnection::MockDatabaseConnection(_) => SchemaBuilderBackend::Postgres,
+            DatabaseConnection::Disconnected => panic!("Disconnected"),
+        }
+    }
+
     pub async fn execute(&self, stmt: Statement) -> Result<ExecResult, ExecErr> {
         match self {
             #[cfg(feature = "sqlx-mysql")]
@@ -139,3 +158,17 @@ impl QueryBuilderBackend {
         .into()
     }
 }
+
+impl SchemaBuilderBackend {
+    pub fn build<S>(&self, statement: &S) -> Statement
+    where
+        S: SchemaStatementBuilder,
+    {
+        match self {
+            Self::MySql => statement.build(MysqlQueryBuilder),
+            Self::Postgres => statement.build(PostgresQueryBuilder),
+            Self::Sqlite => statement.build(SqliteQueryBuilder),
+        }
+        .into()
+    }
+}