From bf595ca9265b1ad07a548ef83977b9253551985b Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Wed, 13 Oct 2021 15:54:49 +0800
Subject: [PATCH] Use Figment config from Rocket

---
 examples/rocket_example/src/main.rs | 20 +++++++-------------
 examples/rocket_example/src/pool.rs | 12 +++++++++++-
 src/database/mod.rs                 | 12 ++++++++----
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/examples/rocket_example/src/main.rs b/examples/rocket_example/src/main.rs
index e2ef9254..2fa77374 100644
--- a/examples/rocket_example/src/main.rs
+++ b/examples/rocket_example/src/main.rs
@@ -48,15 +48,14 @@ async fn create(conn: Connection<'_, Db>, post_form: Form<post::Model>) -> Flash
 }
 
 #[post("/<id>", data = "<post_form>")]
-async fn update(conn: Connection<'_, Db>, id: i32, post_form: Form<post::Model>) -> Flash<Redirect> {
+async fn update(
+    conn: Connection<'_, Db>,
+    id: i32,
+    post_form: Form<post::Model>,
+) -> Flash<Redirect> {
     let db = conn.into_inner();
 
-    let post: post::ActiveModel = Post::find_by_id(id)
-        .one(db)
-        .await
-        .unwrap()
-        .unwrap()
-        .into();
+    let post: post::ActiveModel = Post::find_by_id(id).one(db).await.unwrap().unwrap().into();
 
     let form = post_form.into_inner();
 
@@ -133,12 +132,7 @@ async fn edit(conn: Connection<'_, Db>, id: i32) -> Template {
 async fn delete(conn: Connection<'_, Db>, id: i32) -> Flash<Redirect> {
     let db = conn.into_inner();
 
-    let post: post::ActiveModel = Post::find_by_id(id)
-        .one(db)
-        .await
-        .unwrap()
-        .unwrap()
-        .into();
+    let post: post::ActiveModel = Post::find_by_id(id).one(db).await.unwrap().unwrap().into();
 
     post.delete(db).await.unwrap();
 
diff --git a/examples/rocket_example/src/pool.rs b/examples/rocket_example/src/pool.rs
index afc8d48d..9b729a51 100644
--- a/examples/rocket_example/src/pool.rs
+++ b/examples/rocket_example/src/pool.rs
@@ -1,5 +1,7 @@
 use async_trait::async_trait;
+use sea_orm::ConnectOptions;
 use sea_orm_rocket::{rocket::figment::Figment, Config, Database};
+use std::time::Duration;
 
 #[derive(Database, Debug)]
 #[database("sea_orm")]
@@ -18,7 +20,15 @@ impl sea_orm_rocket::Pool for SeaOrmPool {
 
     async fn init(figment: &Figment) -> Result<Self, Self::Error> {
         let config = figment.extract::<Config>().unwrap();
-        let conn = sea_orm::Database::connect(&config.url).await.unwrap();
+        let mut options: ConnectOptions = config.url.into();
+        options
+            .max_connections(config.max_connections as u32)
+            .min_connections(config.min_connections.unwrap_or_default())
+            .connect_timeout(Duration::from_secs(config.connect_timeout));
+        if let Some(idle_timeout) = config.idle_timeout {
+            options.idle_timeout(Duration::from_secs(idle_timeout));
+        }
+        let conn = sea_orm::Database::connect(options).await.unwrap();
 
         Ok(SeaOrmPool { conn })
     }
diff --git a/src/database/mod.rs b/src/database/mod.rs
index 33a6b7c5..7cb4a57e 100644
--- a/src/database/mod.rs
+++ b/src/database/mod.rs
@@ -115,22 +115,26 @@ impl ConnectOptions {
     }
 
     /// Set the maximum number of connections of the pool
-    pub fn max_connections(&mut self, value: u32) {
+    pub fn max_connections(&mut self, value: u32) -> &mut Self {
         self.max_connections = Some(value);
+        self
     }
 
     /// Set the minimum number of connections of the pool
-    pub fn min_connections(&mut self, value: u32) {
+    pub fn min_connections(&mut self, value: u32) -> &mut Self {
         self.min_connections = Some(value);
+        self
     }
 
     /// Set the timeout duration when acquiring a connection
-    pub fn connect_timeout(&mut self, value: Duration) {
+    pub fn connect_timeout(&mut self, value: Duration) -> &mut Self {
         self.connect_timeout = Some(value);
+        self
     }
 
     /// Set the idle duration before closing a connection
-    pub fn idle_timeout(&mut self, value: Duration) {
+    pub fn idle_timeout(&mut self, value: Duration) -> &mut Self {
         self.idle_timeout = Some(value);
+        self
     }
 }