use crate::health::HealthCheckerTrait;
use crate::models::{DatabaseType, HealthCheckType, HealthResult, Service, ServiceType};
use async_trait::async_trait;
use std::time::{Duration, Instant};
use tracing::trace;
pub struct DatabaseHealthChecker {
timeout: Duration,
}
impl DatabaseHealthChecker {
pub fn new(timeout: Duration) -> Self {
Self { timeout }
}
#[cfg(feature = "db-health-checks")]
async fn check_postgres(&self, host: &str, port: u16, database: Option<&str>) -> HealthResult {
use tokio_postgres::NoTls;
let db = database.unwrap_or("postgres");
let config_str = format!("host={} port={} user=postgres dbname={} connect_timeout={}",
host, port, db, self.timeout.as_secs());
let start = Instant::now();
match tokio_postgres::connect(&config_str, NoTls).await {
Ok((client, connection)) => {
tokio::spawn(async move {
if let Err(e) = connection.await {
trace!("PostgreSQL connection error: {}", e);
}
});
match client.simple_query("SELECT 1").await {
Ok(_) => {
let response_time = start.elapsed().as_millis() as u64;
trace!("PostgreSQL is healthy ({}ms)", response_time);
HealthResult::healthy(response_time)
}
Err(e) => {
HealthResult::unhealthy(
format!("PostgreSQL query failed: {}", e),
None,
)
}
}
}
Err(e) => {
HealthResult::unhealthy(
format!("Cannot connect to PostgreSQL: {}", e),
Some("Check if PostgreSQL is running and accepting connections".to_string()),
)
}
}
}
#[cfg(feature = "db-health-checks")]
async fn check_redis(&self, host: &str, port: u16) -> HealthResult {
let url = format!("redis://{}:{}", host, port);
let start = Instant::now();
match redis::Client::open(url.as_str()) {
Ok(client) => {
match client.get_multiplexed_async_connection().await {
Ok(mut con) => {
match redis::cmd("PING").query_async(&mut con).await {
Ok(response) => {
let response: String = response;
if response == "PONG" {
let response_time = start.elapsed().as_millis() as u64;
trace!("Redis is healthy ({}ms)", response_time);
HealthResult::healthy(response_time)
} else {
HealthResult::unhealthy(
format!("Redis PING returned unexpected response: {}", response),
Some("Redis may be misconfigured".to_string()),
)
}
}
Err(e) => {
HealthResult::unhealthy(
format!("Redis PING failed: {}", e),
Some("Check if Redis is running and accepting connections".to_string()),
)
}
}
}
Err(e) => {
HealthResult::unhealthy(
format!("Cannot connect to Redis: {}", e),
Some("Check if Redis is running and accepting connections".to_string()),
)
}
}
}
Err(e) => {
HealthResult::unhealthy(
format!("Invalid Redis URL: {}", e),
None,
)
}
}
}
}
#[async_trait]
impl HealthCheckerTrait for DatabaseHealthChecker {
async fn check(&self, service: &Service) -> HealthResult {
match &service.health_check.check_type {
#[cfg(feature = "db-health-checks")]
HealthCheckType::Postgres { database } => {
self.check_postgres(&service.host, service.port, database.as_deref()).await
}
#[cfg(feature = "db-health-checks")]
HealthCheckType::Redis => {
self.check_redis(&service.host, service.port).await
}
_ => {
match &service.service_type {
#[cfg(feature = "db-health-checks")]
ServiceType::Database(DatabaseType::Postgres) => {
self.check_postgres(&service.host, service.port, None).await
}
#[cfg(feature = "db-health-checks")]
ServiceType::Cache(crate::models::CacheType::Redis) => {
self.check_redis(&service.host, service.port).await
}
_ => {
HealthResult::unhealthy(
"Database health check not implemented for this type".to_string(),
Some("Use port connectivity check instead".to_string()),
)
}
}
}
}
}
fn supports(&self, service: &Service) -> bool {
matches!(service.health_check.check_type, HealthCheckType::Postgres { .. })
|| matches!(service.health_check.check_type, HealthCheckType::Redis)
|| matches!(service.service_type, ServiceType::Database(_))
|| matches!(service.service_type, ServiceType::Cache(_))
}
}