b453d5d
use std::sync::Arc;
use axum::Router;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use sea_orm::{ConnectionTrait, DbBackend, Statement};
use crate::{AppDatabase, AppRedis, AppState};
#[derive(Clone)]
struct HealthState {
database: Option<Arc<AppDatabase>>,
redis: AppRedis,
kenauth: Option<Arc<kenauth::KenAuthClient>>,
}
async fn healthz() -> Response {
(StatusCode::OK, "ok").into_response()
}
async fn readyz(State(state): State<HealthState>) -> Response {
let db_ok = match state.database.as_ref() {
None => true,
Some(db) => db
.read()
.query_one_raw(Statement::from_sql_and_values(
DbBackend::Postgres,
"SELECT 1",
[],
))
.await
.is_ok(),
};
let mut conn = state.redis.connection_clone();
let redis_ok: bool = match &mut conn {
crate::redis::RedisConnection::Standalone(c) => {
redis::cmd("PING").query_async::<String>(c).await.is_ok()
}
crate::redis::RedisConnection::Cluster(c) => {
redis::cmd("PING").query_async::<String>(c).await.is_ok()
}
};
let kenauth_ok = match state.kenauth.as_ref() {
None => true,
Some(client) => client.check_health().await.is_ok(),
};
if db_ok && redis_ok && kenauth_ok {
(StatusCode::OK, "ready").into_response()
} else {
(StatusCode::SERVICE_UNAVAILABLE, "not ready").into_response()
}
}
pub fn router_with(
database: Option<Arc<AppDatabase>>,
redis: AppRedis,
kenauth: Option<Arc<kenauth::KenAuthClient>>,
) -> Router {
Router::new()
.route("/healthz", axum::routing::get(healthz))
.route("/readyz", axum::routing::get(readyz))
.with_state(HealthState {
database,
redis,
kenauth,
})
}
pub fn router(state: AppState) -> Router {
router_with(
Some(state.database.clone()),
state.redis.clone(),
state.kenauth.clone(),
)
}