1 files changed, 26 insertions, 20 deletions

diff --git a/src/main.rs b/src/main.rs
index 10b3cc7..5462870 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,16 +8,18 @@ pub mod routes;
 #[macro_use]
 pub mod schema;
 
-use actix_rt;
-use actix_web::{middleware, App, HttpServer};
+use actix_web::{middleware, web, App, HttpServer};
 use diesel::pg::PgConnection;
-use diesel::prelude::*;
-use dotenv::dotenv;
+use diesel::r2d2;
 use std::{env, io::Error};
 
-#[cfg_attr(tarpaulin, skip)]
-fn main() {
-    if let Err(ref e) = run() {
+/// Short-hand for the database pool type to use throughout the app.
+type DbPool = r2d2::Pool<r2d2::ConnectionManager<PgConnection>>;
+
+#[cfg(not(tarpaulin_include))]
+#[actix_web::main]
+async fn main() {
+    if let Err(ref e) = run().await {
         error!("error: {}", e);
 
         ::std::process::exit(1);
@@ -25,27 +27,31 @@ fn main() {
 }
 
 // Run start-up for the server and dependencies
-#[cfg_attr(tarpaulin, skip)]
-fn run() -> Result<(), Error> {
-    dotenv().ok();
-    let db_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
+#[cfg(not(tarpaulin_include))]
+async fn run() -> Result<(), Error> {
+    dotenvy::dotenv().ok();
     let bind_address = env::var("BIND_ADDRESS").expect("BIND_ADDRESS must be set");
 
     env_logger::init();
-    let sys = actix_rt::System::new("august-offensive");
+    let pool = init_db_pool();
 
-    PgConnection::establish(&db_url).expect(&format!("Error connecting to {}", db_url));
-
-    HttpServer::new(|| {
+    info!("Started http server: {}", bind_address);
+    HttpServer::new(move || {
         App::new()
+            .app_data(web::Data::new(pool.clone()))
             .wrap(middleware::Logger::default())
             .service(routes::get_scope())
     })
     .bind(&bind_address)?
-    .start();
-
-    info!("Started http server: {}", bind_address);
+    .run()
+    .await
+}
 
-    sys.run()?;
-    Ok(())
+#[cfg(not(tarpaulin_include))]
+fn init_db_pool() -> DbPool {
+    let db_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
+    let manager = r2d2::ConnectionManager::<PgConnection>::new(db_url);
+    r2d2::Pool::builder()
+        .build(manager)
+        .expect("Database URL is not valid")
 }