From e4ddfbd025f433da79387bee5707ab369d48b298 Mon Sep 17 00:00:00 2001
From: "Dustin C. Hatch" <dustin@hatch.name>
Date: Sun, 9 Mar 2025 20:55:26 -0500
Subject: [PATCH] Run DB migrations at startup

Naturally, we need the database schema in place in order to use it.
---
 .containerignore |  1 +
 Cargo.toml       |  2 +-
 Containerfile    |  1 +
 src/main.rs      | 20 ++++++++++++++++++--
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/.containerignore b/.containerignore
index bc6f674..dee510c 100644
--- a/.containerignore
+++ b/.containerignore
@@ -2,5 +2,6 @@
 !Cargo.*
 !.sqlx/
 !js/
+!migrations/
 !src/
 !templates/
diff --git a/Cargo.toml b/Cargo.toml
index 7880e0c..2273fee 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,7 +17,7 @@ rocket_db_pools = { version = "0.2.0", features = ["sqlx_macros", "sqlx_postgres
 rocket_dyn_templates = { version = "0.2.0", features = ["tera"] }
 rust_decimal = { version = "1.36.0", features = ["serde-with-str"] }
 serde = { version = "1.0.218", default-features = false, features = ["derive"] }
-sqlx = { version = "~0.7.4", default-features = false, features = ["chrono", "macros", "postgres", "rust_decimal", "time"] }
+sqlx = { version = "~0.7.4", default-features = false, features = ["chrono", "macros", "migrate", "postgres", "rust_decimal", "time"] }
 thiserror = "2.0.12"
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
diff --git a/Containerfile b/Containerfile
index f368222..ccd7a6b 100644
--- a/Containerfile
+++ b/Containerfile
@@ -13,6 +13,7 @@ WORKDIR /build
 COPY Cargo.* .
 COPY src src
 COPY .sqlx .sqlx
+COPY migrations migrations
 
 RUN --mount=type=cache,target=/root/.cargo \
 	cargo build --release --locked
diff --git a/src/main.rs b/src/main.rs
index ddb93bd..c4dd60d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,16 +2,17 @@ mod config;
 mod firefly;
 mod receipts;
 
+use rocket::fairing::{self, AdHoc};
 use rocket::form::Form;
 use rocket::fs::{FileServer, TempFile};
 use rocket::http::Status;
 use rocket::response::Redirect;
 use rocket::tokio::io::{AsyncReadExt, BufReader};
-use rocket::State;
+use rocket::{Rocket, State};
 use rocket_db_pools::Database as RocketDatabase;
 use rocket_dyn_templates::{context, Template};
 use serde::Serialize;
-use tracing::{debug, error};
+use tracing::{debug, error, info};
 
 use config::Config;
 use firefly::{
@@ -229,6 +230,20 @@ async fn update_transaction(
     (Status::Ok, "Successfully updated transaction".into())
 }
 
+async fn run_migrations(rocket: Rocket<rocket::Build>) -> fairing::Result {
+    if let Some(db) = Database::fetch(&rocket) {
+        info!("Applying database migrations");
+        if let Err(e) = sqlx::migrate!("./migrations").run(&db.0).await {
+            error!("Database migration failed: {}", e);
+            Err(rocket)
+        } else {
+            Ok(rocket)
+        }
+    } else {
+        Err(rocket)
+    }
+}
+
 #[rocket::launch]
 async fn rocket() -> _ {
     tracing_subscriber::fmt()
@@ -265,4 +280,5 @@ async fn rocket() -> _ {
         .mount("/static", FileServer::from("static"))
         .attach(Template::fairing())
         .attach(Database::init())
+        .attach(AdHoc::try_on_ignite("Migrate Database", run_migrations))
 }