diff --git a/src/main.rs b/src/main.rs index e31439b..79892dd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,8 +9,10 @@ mod feeds; mod poll; mod user; +use rocket::fairing::{self, AdHoc}; use rocket::fs::FileServer; use rocket::response::Redirect; +use rocket::{Build, Rocket}; use rocket_db_pools::{sqlx, Connection, Database}; use rocket_dyn_templates::{context, Template}; use user::AuthenticatedUser; @@ -59,6 +61,26 @@ fn login() -> Template { Template::render("login", context! {}) } +// Run migrations and setup demo data if needed +async fn setup_database(demo: bool, rocket: Rocket) -> fairing::Result { + let db = match Db::fetch(&rocket) { + Some(db) => db, + None => return Err(rocket), + }; + + let pool = db.0.clone(); + sqlx::migrate!("./migrations") + .run(&pool) + .await + .expect("Failed to run database migrations"); + + if demo { + demo::setup_demo_data(&pool).await; + } + + Ok(rocket) +} + #[launch] fn rocket() -> _ { let args = Args::parse(); @@ -77,21 +99,8 @@ fn rocket() -> _ { format!("sqlite:{}", database) }; - // Run migrations and setup demo data if needed - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async { - let pool = sqlx::SqlitePool::connect(&db_url).await.unwrap(); - sqlx::migrate!("./migrations") - .run(&pool) - .await - .expect("Failed to run database migrations"); - - if args.demo { - demo::setup_demo_data(&pool).await; - } - }); - let figment = rocket::Config::figment().merge(("databases.rss_data.url", db_url)); + let demo = args.demo; rocket::custom(figment) .mount( @@ -117,4 +126,7 @@ fn rocket() -> _ { .mount("/static", FileServer::from("static")) .attach(Template::fairing()) .attach(Db::init()) + .attach(AdHoc::try_on_ignite("DB Setup", move |rocket| async move { + setup_database(demo, rocket).await + })) }