From 7fb3565992580e307c70a2012f19511aeec694e6 Mon Sep 17 00:00:00 2001
From: Takayuki Maeda <41065217+TaKO8Ki@users.noreply.github.com>
Date: Fri, 10 Sep 2021 01:13:21 +0900
Subject: [PATCH] use pool options (#86)

---
 src/database/mysql.rs    | 10 +++++++---
 src/database/postgres.rs |  8 ++++++--
 src/database/sqlite.rs   | 10 +++++++---
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/database/mysql.rs b/src/database/mysql.rs
index 8241d97..16d7b6c 100644
--- a/src/database/mysql.rs
+++ b/src/database/mysql.rs
@@ -3,17 +3,21 @@ use async_trait::async_trait;
 use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
 use database_tree::{Child, Database, Table};
 use futures::TryStreamExt;
-use sqlx::mysql::{MySqlColumn, MySqlPool as MPool, MySqlRow};
+use sqlx::mysql::{MySqlColumn, MySqlPoolOptions, MySqlRow};
 use sqlx::{Column as _, Row as _, TypeInfo as _};
+use std::time::Duration;
 
 pub struct MySqlPool {
-    pool: MPool,
+    pool: sqlx::mysql::MySqlPool,
 }
 
 impl MySqlPool {
     pub async fn new(database_url: &str) -> anyhow::Result<Self> {
         Ok(Self {
-            pool: MPool::connect(database_url).await?,
+            pool: MySqlPoolOptions::new()
+                .connect_timeout(Duration::from_millis(500))
+                .connect(database_url)
+                .await?,
         })
     }
 }
diff --git a/src/database/postgres.rs b/src/database/postgres.rs
index f600c06..f8219a8 100644
--- a/src/database/postgres.rs
+++ b/src/database/postgres.rs
@@ -4,8 +4,9 @@ use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
 use database_tree::{Child, Database, Schema, Table};
 use futures::TryStreamExt;
 use itertools::Itertools;
-use sqlx::postgres::{PgColumn, PgPool, PgRow};
+use sqlx::postgres::{PgColumn, PgPool, PgPoolOptions, PgRow};
 use sqlx::{Column as _, Row as _, TypeInfo as _};
+use std::time::Duration;
 
 pub struct PostgresPool {
     pool: PgPool,
@@ -14,7 +15,10 @@ pub struct PostgresPool {
 impl PostgresPool {
     pub async fn new(database_url: &str) -> anyhow::Result<Self> {
         Ok(Self {
-            pool: PgPool::connect(database_url).await?,
+            pool: PgPoolOptions::new()
+                .connect_timeout(Duration::from_millis(500))
+                .connect(database_url)
+                .await?,
         })
     }
 }
diff --git a/src/database/sqlite.rs b/src/database/sqlite.rs
index 26d1771..6ca8e98 100644
--- a/src/database/sqlite.rs
+++ b/src/database/sqlite.rs
@@ -3,17 +3,21 @@ use async_trait::async_trait;
 use chrono::NaiveDateTime;
 use database_tree::{Child, Database, Table};
 use futures::TryStreamExt;
-use sqlx::sqlite::{SqliteColumn, SqlitePool as SPool, SqliteRow};
+use sqlx::sqlite::{SqliteColumn, SqlitePoolOptions, SqliteRow};
 use sqlx::{Column as _, Row as _, TypeInfo as _};
+use std::time::Duration;
 
 pub struct SqlitePool {
-    pool: SPool,
+    pool: sqlx::sqlite::SqlitePool,
 }
 
 impl SqlitePool {
     pub async fn new(database_url: &str) -> anyhow::Result<Self> {
         Ok(Self {
-            pool: SPool::connect(database_url).await?,
+            pool: SqlitePoolOptions::new()
+                .connect_timeout(Duration::from_millis(500))
+                .connect(database_url)
+                .await?,
         })
     }
 }