Allow database timeouts to be configured

src/app.rs

@@ -3,6 +3,7 @@
 use std::env;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex};
+use std::time::Duration;
 
 use diesel::r2d2;
 use git2;
@@ -71,11 +72,19 @@ impl App {
             _ => 1,
         };
 
+        let db_connection_timeout = match (env::var("DB_TIMEOUT"), config.env) {
+            (Ok(num), _) => num.parse().expect("couldn't parse DB_TIMEOUT"),
+            (_, ::Env::Production) => 10,
+            _ => 30,
+        };
+
         let thread_pool = Arc::new(ScheduledThreadPool::new(db_helper_threads));
 
         let diesel_db_config = r2d2::Pool::builder()
             .max_size(db_pool_size)
             .min_idle(db_min_idle)
+            .connection_timeout(Duration::from_secs(db_connection_timeout))
+            .connection_customizer(Box::new(db::SetStatementTimeout(db_connection_timeout)))
             .thread_pool(thread_pool);
 
         let repo = git2::Repository::open(&config.git_repo_checkout).unwrap();

src/db.rs

@@ -1,8 +1,8 @@
 use std::env;
 
 use conduit::Request;
-use diesel::prelude::{ConnectionResult, PgConnection};
-use diesel::r2d2::{self, ConnectionManager};
+use diesel::prelude::*;
+use diesel::r2d2::{self, ConnectionManager, CustomizeConnection};
 use url::Url;
 
 use middleware::app::RequestApp;
@@ -44,3 +44,17 @@ impl<T: Request + ?Sized> RequestTransaction for T {
         self.app().diesel_database.get().map_err(Into::into)
     }
 }
+
+#[derive(Debug, Clone, Copy)]
+pub struct SetStatementTimeout(pub u64);
+
+impl CustomizeConnection<PgConnection, r2d2::Error> for SetStatementTimeout {
+    fn on_acquire(&self, conn: &mut PgConnection) -> Result<(), r2d2::Error> {
+        use diesel::sql_query;
+
+        sql_query(format!("SET statement_timeout = {}", self.0 * 1000))
+            .execute(conn)
+            .map_err(r2d2::Error::QueryError)?;
+        Ok(())
+    }
+}