diff --git a/kepler-core/src/db.rs b/kepler-core/src/db.rs index fd99fe58..5faa414c 100644 --- a/kepler-core/src/db.rs +++ b/kepler-core/src/db.rs @@ -142,6 +142,17 @@ where } } +impl OrbitDatabase +where + C: TransactionTrait, +{ + pub async fn check_db_connection(&self) -> Result<(), DbErr> { + // there's a `ping` method on the connection, but we can't access it from here + // but starting a transaction should be enough to check the connection + self.conn.begin().await.map(|_| ()) + } +} + pub type InvocationInputs = HashMap<(OrbitId, String), (Metadata, HashBuffer)>; impl OrbitDatabase diff --git a/src/lib.rs b/src/lib.rs index 9dce2390..a348b350 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,7 +22,7 @@ mod tracing; use config::{BlockStorage, Config, Keys, StagingStorage}; use kepler_core::{ keys::{SecretsSetup, StaticSecret}, - sea_orm::{Database, DatabaseConnection}, + sea_orm::{ConnectOptions, Database, DatabaseConnection}, storage::{either::Either, memory::MemoryStaging, StorageConfig}, OrbitDatabase, }; @@ -86,8 +86,11 @@ pub async fn app(config: &Figment) -> Result> { Keys::Static(s) => s.try_into()?, }; + let mut connect_opts = ConnectOptions::from(&kepler_config.storage.database); + connect_opts.max_connections(100); + let kepler = Kepler::new( - Database::connect(&kepler_config.storage.database).await?, + Database::connect(connect_opts).await?, kepler_config.storage.blocks.open().await?, key_setup.setup(()).await?, ) diff --git a/src/routes/mod.rs b/src/routes/mod.rs index 5aa4f12f..29bf536c 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -12,10 +12,11 @@ use crate::{ BlockStage, BlockStores, Kepler, }; use kepler_core::{ + sea_orm::DbErr, storage::{ImmutableReadStore, ImmutableStaging}, types::Resource, util::{DelegationInfo, InvocationInfo}, - TxError, + TxError, TxStoreError, }; pub mod util; @@ -23,11 +24,19 @@ use util::LimitedReader; #[allow(clippy::let_unit_value)] pub mod util_routes { + use super::*; + #[options("/<_s..>")] pub async fn cors(_s: std::path::PathBuf) {} #[get("/healthz")] - pub fn healthcheck() {} + pub async fn healthcheck(s: &State) -> Status { + if s.check_db_connection().await.is_ok() { + Status::Ok + } else { + Status::InternalServerError + } + } } #[get("/peer/generate/")] @@ -69,6 +78,7 @@ pub async fn delegate( ( match e { TxError::OrbitNotFound => Status::NotFound, + TxError::Db(DbErr::ConnectionAcquire) => Status::InternalServerError, _ => Status::Unauthorized, }, e.to_string(), @@ -183,7 +193,18 @@ pub async fn invoke( _ => unreachable!(), }, ) - .map_err(|e| (Status::Unauthorized, e.to_string())); + .map_err(|e| { + ( + match e { + TxStoreError::Tx(TxError::OrbitNotFound) => Status::NotFound, + TxStoreError::Tx(TxError::Db(DbErr::ConnectionAcquire)) => { + Status::InternalServerError + } + _ => Status::Unauthorized, + }, + e.to_string(), + ) + }); timer.observe_duration(); res