Retry initial db connection, with adjustable option

src/config.rs

@@ -347,6 +347,9 @@ make_config! {
         /// that do not support WAL. Please make sure you read project wiki on the topic before changing this setting.
         enable_db_wal:          bool,   false,  def,    true;
 
+        /// Max database connection retries |> Number of times to retry the database connection during startup, with 1 second between each retry, set to 0 to retry indefinitely
+        db_connection_retries:  u32,    false,  def,    15;
+
         /// Bypass admin page security (Know the risks!) |> Disables the Admin Token for the admin page so you may use your own auth in-front
         disable_admin_token:    bool,   true,   def,    false;
 

src/db/mod.rs

@@ -49,7 +49,7 @@ macro_rules! generate_connections {
                     DbConnType::$name => {
                         #[cfg($name)]
                         {
-                            paste::paste!{ [< $name _migrations >]::run_migrations(); }
+                            paste::paste!{ [< $name _migrations >]::run_migrations()?; }
                             let manager = ConnectionManager::new(&url);
                             let pool = Pool::builder().build(manager).map_res("Failed to create pool")?;
                             return Ok(Self::$name(pool));
@@ -242,7 +242,7 @@ mod sqlite_migrations {
     #[allow(unused_imports)]
     embed_migrations!("migrations/sqlite");
 
-    pub fn run_migrations() {
+    pub fn run_migrations() -> Result<(), super::Error> {
         // Make sure the directory exists
         let url = crate::CONFIG.database_url();
         let path = std::path::Path::new(&url);
@@ -257,7 +257,7 @@ mod sqlite_migrations {
         use diesel::{Connection, RunQueryDsl};
         // Make sure the database is up to date (create if it doesn't exist, or run the migrations)
         let connection =
-            diesel::sqlite::SqliteConnection::establish(&crate::CONFIG.database_url()).expect("Can't connect to DB");
+            diesel::sqlite::SqliteConnection::establish(&crate::CONFIG.database_url())?;
         // Disable Foreign Key Checks during migration
         
         // Scoped to a connection.
@@ -272,7 +272,8 @@ mod sqlite_migrations {
                 .expect("Failed to turn on WAL");
         }
 
-        embedded_migrations::run_with_output(&connection, &mut std::io::stdout()).expect("Can't run migrations");
+        embedded_migrations::run_with_output(&connection, &mut std::io::stdout())?;
+        Ok(())
     }
 }
 
@@ -281,11 +282,11 @@ mod mysql_migrations {
     #[allow(unused_imports)]
     embed_migrations!("migrations/mysql");
 
-    pub fn run_migrations() {
+    pub fn run_migrations() -> Result<(), super::Error> {
         use diesel::{Connection, RunQueryDsl};
         // Make sure the database is up to date (create if it doesn't exist, or run the migrations)
         let connection =
-            diesel::mysql::MysqlConnection::establish(&crate::CONFIG.database_url()).expect("Can't connect to DB");
+            diesel::mysql::MysqlConnection::establish(&crate::CONFIG.database_url())?;
         // Disable Foreign Key Checks during migration
 
         // Scoped to a connection/session.
@@ -293,7 +294,8 @@ mod mysql_migrations {
             .execute(&connection)
             .expect("Failed to disable Foreign Key Checks during migrations");
 
-        embedded_migrations::run_with_output(&connection, &mut std::io::stdout()).expect("Can't run migrations");
+        embedded_migrations::run_with_output(&connection, &mut std::io::stdout())?;
+        Ok(())
     }
 }
 
@@ -302,11 +304,11 @@ mod postgresql_migrations {
     #[allow(unused_imports)]
     embed_migrations!("migrations/postgresql");
 
-    pub fn run_migrations() {
+    pub fn run_migrations() -> Result<(), super::Error> {
         use diesel::{Connection, RunQueryDsl};
         // Make sure the database is up to date (create if it doesn't exist, or run the migrations)
         let connection =
-            diesel::pg::PgConnection::establish(&crate::CONFIG.database_url()).expect("Can't connect to DB");
+            diesel::pg::PgConnection::establish(&crate::CONFIG.database_url())?;
         // Disable Foreign Key Checks during migration
         
         // FIXME: Per https://www.postgresql.org/docs/12/sql-set-constraints.html,
@@ -319,6 +321,7 @@ mod postgresql_migrations {
             .execute(&connection)
             .expect("Failed to disable Foreign Key Checks during migrations");
 
-        embedded_migrations::run_with_output(&connection, &mut std::io::stdout()).expect("Can't run migrations");
+        embedded_migrations::run_with_output(&connection, &mut std::io::stdout())?;
+        Ok(())
     }
 }

src/error.rs

@@ -34,6 +34,8 @@ macro_rules! make_error {
 }
 
 use diesel::result::Error as DieselErr;
+use diesel::ConnectionError as DieselConErr;
+use diesel_migrations::RunMigrationsError as DieselMigErr;
 use diesel::r2d2::PoolError as R2d2Err;
 use handlebars::RenderError as HbErr;
 use jsonwebtoken::errors::Error as JWTErr;
@@ -83,6 +85,9 @@ make_error! {
     AddressError(AddrErr):    _has_source, _api_error,
     SmtpError(SmtpErr):       _has_source, _api_error,
     FromStrError(FromStrErr): _has_source, _api_error,
+
+    DieselConError(DieselConErr): _has_source, _api_error,
+    DieselMigError(DieselMigErr): _has_source, _api_error,
 }
 
 impl std::fmt::Debug for Error {

src/main.rs

@@ -268,7 +268,7 @@ fn check_web_vault() {
 }
 
 fn launch_rocket(extra_debug: bool) {
-    let pool = match db::DbPool::from_config() {
+    let pool = match util::retry_db(db::DbPool::from_config, CONFIG.db_connection_retries()) {
         Ok(p) => p,
         Err(e) => {
             error!("Error creating database pool: {:?}", e);

src/util.rs

@@ -410,7 +410,7 @@ fn _process_key(key: &str) -> String {
 // Retry methods
 //
 
-pub fn retry<F, T, E>(func: F, max_tries: i32) -> Result<T, E>
+pub fn retry<F, T, E>(func: F, max_tries: u32) -> Result<T, E>
 where
     F: Fn() -> Result<T, E>,
 {
@@ -432,3 +432,29 @@ where
         }
     }
 }
+
+pub fn retry_db<F, T, E>(func: F, max_tries: u32) -> Result<T, E>
+where
+    F: Fn() -> Result<T, E>,
+    E: std::error::Error,
+{
+    use std::{thread::sleep, time::Duration};
+    let mut tries = 0;
+
+    loop {
+        match func() {
+            ok @ Ok(_) => return ok,
+            Err(e) => {
+                tries += 1;
+
+                if tries >= max_tries && max_tries > 0 {
+                    return Err(e);
+                }
+
+                warn!("Can't connect to database, retrying: {:?}", e);
+
+                sleep(Duration::from_millis(1_000));
+            }
+        }
+    }
+}