0
0
Fork 0
mirror of https://github.com/dani-garcia/vaultwarden synced 2024-12-13 17:22:58 +01:00

Retry initial db connection, with adjustable option

This commit is contained in:
Daniel García 2020-10-03 22:31:52 +02:00
parent 22b9c80007
commit 729c9cff41
No known key found for this signature in database
GPG key ID: FC8A7D14C3CD543A
5 changed files with 49 additions and 12 deletions

View file

@ -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;

View file

@ -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(())
}
}

View file

@ -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 {

View file

@ -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);

View file

@ -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));
}
}
}
}