Fix issue with demo data setup

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<Build>) -> 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
+        }))
 }