From dd4b04282c7fa6fc569c35a841e41f581d247c27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Emil=20Priv=C3=A9r?= <emil@priver.dev>
Date: Fri, 22 Dec 2023 09:13:18 +0100
Subject: [PATCH] added waiting

---
 src/config.rs                    | 10 ++++++++++
 src/database_drivers/libsql.rs   |  1 +
 src/database_drivers/maria.rs    | 20 +++++++++++++++++++-
 src/database_drivers/mysql.rs    | 19 ++++++++++++++++++-
 src/database_drivers/postgres.rs | 19 ++++++++++++++++++-
 5 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index db742ce..9b0e79f 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -25,6 +25,16 @@ pub fn database_token() -> Option<String> {
     None
 }
 
+pub fn wait_timeout() -> usize {
+    if let Ok(v) = env::var("DATABASE_WAIT_TIMEOUT") {
+        if let Ok(v) = v.parse::<usize>() {
+            return v;
+        }
+    }
+
+    0
+}
+
 pub enum Database {
     LibSQL,
     Postgres,
diff --git a/src/database_drivers/libsql.rs b/src/database_drivers/libsql.rs
index e7b1ac2..45df8ed 100644
--- a/src/database_drivers/libsql.rs
+++ b/src/database_drivers/libsql.rs
@@ -21,6 +21,7 @@ impl<'a> LibSQLDriver {
             Ok(c) => c,
             Err(err) => bail!("{:?}", err),
         };
+
         Ok(LibSQLDriver { db: client })
     }
 }
diff --git a/src/database_drivers/maria.rs b/src/database_drivers/maria.rs
index 70efe9d..dc49b72 100644
--- a/src/database_drivers/maria.rs
+++ b/src/database_drivers/maria.rs
@@ -1,3 +1,4 @@
+use crate::config;
 use crate::database_drivers::DatabaseDriver;
 use anyhow::Result;
 use sqlx::mysql::MySqlRow;
@@ -11,7 +12,24 @@ pub struct MariaDBDriver {
 
 impl<'a> MariaDBDriver {
     pub async fn new<'b>(db_url: &str) -> Result<MariaDBDriver> {
-        let client = MySqlConnection::connect(db_url).await?;
+        let mut client = MySqlConnection::connect(db_url).await?;
+
+        let wait_timeout = config::wait_timeout();
+
+        let mut count = 0;
+        loop {
+            if count > wait_timeout {
+                break;
+            }
+
+            if client.ping().await.is_ok() {
+                break;
+            }
+
+            count += 1;
+            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+        }
+
         Ok(MariaDBDriver { db: client })
     }
 }
diff --git a/src/database_drivers/mysql.rs b/src/database_drivers/mysql.rs
index ce216be..463104c 100644
--- a/src/database_drivers/mysql.rs
+++ b/src/database_drivers/mysql.rs
@@ -1,3 +1,4 @@
+use crate::config;
 use crate::database_drivers::DatabaseDriver;
 use anyhow::Result;
 use sqlx::mysql::MySqlRow;
@@ -11,7 +12,23 @@ pub struct MySQLDriver {
 
 impl<'a> MySQLDriver {
     pub async fn new<'b>(db_url: &str) -> Result<MySQLDriver> {
-        let client = MySqlConnection::connect(db_url).await?;
+        let mut client = MySqlConnection::connect(db_url).await?;
+
+        let wait_timeout = config::wait_timeout();
+        let mut count = 0;
+        loop {
+            if count > wait_timeout {
+                break;
+            }
+
+            if client.ping().await.is_ok() {
+                break;
+            }
+
+            count += 1;
+            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+        }
+
         Ok(MySQLDriver { db: client })
     }
 }
diff --git a/src/database_drivers/postgres.rs b/src/database_drivers/postgres.rs
index 350c0a5..74622e1 100644
--- a/src/database_drivers/postgres.rs
+++ b/src/database_drivers/postgres.rs
@@ -1,3 +1,4 @@
+use crate::config;
 use crate::database_drivers::DatabaseDriver;
 use anyhow::Result;
 use sqlx::postgres::PgRow;
@@ -11,7 +12,23 @@ pub struct PostgresDriver {
 
 impl<'a> PostgresDriver {
     pub async fn new<'b>(db_url: &str) -> Result<PostgresDriver> {
-        let client = PgConnection::connect(db_url).await?;
+        let mut client = PgConnection::connect(db_url).await?;
+
+        let wait_timeout = config::wait_timeout();
+        let mut count = 0;
+        loop {
+            if count > wait_timeout {
+                break;
+            }
+
+            if client.ping().await.is_ok() {
+                break;
+            }
+
+            count += 1;
+            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+        }
+
         Ok(PostgresDriver { db: client })
     }
 }