Add DB ping to healthz

EPFL-ENAC · Oct 23, 2024 · d9042a3 · d9042a3
1 parent 480fd94
commit d9042a3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/src/common/views.rs b/src/common/views.rs
@@ -1,7 +1,8 @@
 use super::models::HealthCheck;
-use crate::common::models::UIConfiguration;
 use crate::external::k8s::services::get_pods;
-use axum::{http::StatusCode, Json};
+use crate::{common::models::UIConfiguration, external::db};
+use axum::{extract::State, http::StatusCode, Json};
+use sea_orm::DatabaseConnection;
 
 #[utoipa::path(
     get,
@@ -15,17 +16,21 @@ use axum::{http::StatusCode, Json};
         )
     )
 )]
-pub async fn healthz() -> (StatusCode, Json<HealthCheck>) {
+pub async fn healthz(State(db): State<DatabaseConnection>) -> (StatusCode, Json<HealthCheck>) {
     // Get health of the API.
     match get_pods().await {
-        Ok(_) => {
+        Err(_) => {
             return (
-                StatusCode::OK,
+                StatusCode::INTERNAL_SERVER_ERROR,
                 Json(HealthCheck {
-                    status: "ok".to_string(),
+                    status: "error".to_string(),
                 }),
             )
         }
+        _ => {}
+    };
+
+    match db.ping().await {
         Err(_) => {
             return (
                 StatusCode::INTERNAL_SERVER_ERROR,
@@ -34,7 +39,15 @@ pub async fn healthz() -> (StatusCode, Json<HealthCheck>) {
                 }),
             )
         }
+        _ => {}
     };
+
+    (
+        StatusCode::OK,
+        Json(HealthCheck {
+            status: "ok".to_string(),
+        }),
+    )
 }
 
 #[utoipa::path(

diff --git a/src/main.rs b/src/main.rs
@@ -47,6 +47,8 @@ async fn main() {
 
     // Set up your Axum app
     let app: Router = Router::new()
+        .route("/healthz", get(common::views::healthz))
+        .route("/api/config", get(common::views::get_ui_config))
         .with_state(db.clone())
         .nest(
             "/api/submissions",
@@ -59,9 +61,7 @@ async fn main() {
         .nest(
             "/tus",
             external::tus::views::router(db.clone(), keycloak_auth_instance),
-        )
-        .route("/healthz", get(common::views::healthz))
-        .route("/api/config", get(common::views::get_ui_config));
+        );
 
     let addr: std::net::SocketAddr = "0.0.0.0:3000".parse().unwrap();
     println!("Listening on {}", addr);