diff --git a/regtest b/regtest index 6e103c7..0f9ea7b 160000 --- a/regtest +++ b/regtest @@ -1 +1 @@ -Subproject commit 6e103c773c3fbe35ff341e56b9c5a3893cb2681b +Subproject commit 0f9ea7bf042213f1b638cb2f46bd981a6fb25376 diff --git a/src/database/mod.rs b/src/database/mod.rs index 19bad1a..c20e8b3 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -1,9 +1,11 @@ +use diesel::connection::SimpleConnection; use diesel::prelude::*; use diesel::r2d2::ConnectionManager; use diesel::{r2d2, PgConnection, SqliteConnection}; use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use log::{debug, info, trace}; use std::error::Error; +use std::time::Duration; pub mod helpers; pub mod model; @@ -19,6 +21,28 @@ pub enum AnyConnection { Sqlite(SqliteConnection), } +#[derive(Debug)] +pub struct ConnectionOptions { + pub busy_timeout: Option, +} + +impl diesel::r2d2::CustomizeConnection for ConnectionOptions { + fn on_acquire(&self, conn: &mut AnyConnection) -> Result<(), diesel::r2d2::Error> { + (|| { + match conn { + AnyConnection::Sqlite(conn) => { + if let Some(d) = self.busy_timeout { + conn.batch_execute(&format!("PRAGMA busy_timeout = {};", d.as_millis()))?; + } + } + _ => {} + } + Ok(()) + })() + .map_err(diesel::r2d2::Error::QueryError) + } +} + pub type Pool = r2d2::Pool>; pub fn connect(url: &str) -> Result> { @@ -30,7 +54,11 @@ pub fn connect(url: &str) -> Result> { debug!("Connecting to {} database", db_name); let manager: ConnectionManager = ConnectionManager::new(url); - let pool = Pool::builder().build(manager)?; + let pool = Pool::builder() + .connection_customizer(Box::new(ConnectionOptions { + busy_timeout: Some(Duration::from_secs(5)), + })) + .build(manager)?; info!("Connected to {} database", db_name);