use std::time::Duration;
use async_trait::async_trait;
use tokio_postgres::Config;
use super::FailureCategory;
#[derive(Debug)]
pub struct HealthFailure {
category: FailureCategory,
detail: String,
}
impl HealthFailure {
pub fn new(category: FailureCategory, detail: impl Into<String>) -> Self {
Self {
category,
detail: detail.into(),
}
}
pub fn unavailable(detail: impl Into<String>) -> Self {
Self::new(FailureCategory::Unavailable, detail)
}
pub fn category(&self) -> FailureCategory {
self.category
}
pub fn detail(&self) -> &str {
&self.detail
}
}
#[async_trait]
pub trait BackendHealth: Send + Sync {
fn backend(&self) -> &'static str;
fn bound(&self) -> Duration;
async fn check(&self) -> Result<(), HealthFailure>;
}
pub struct PostgresHealth {
backend: &'static str,
config: Config,
bound: Duration,
}
impl PostgresHealth {
pub fn new(backend: &'static str, config: Config, bound: Duration) -> Self {
Self {
backend,
config,
bound,
}
}
}
#[async_trait]
impl BackendHealth for PostgresHealth {
fn backend(&self) -> &'static str {
self.backend
}
fn bound(&self) -> Duration {
self.bound
}
async fn check(&self) -> Result<(), HealthFailure> {
let (client, connection) = self
.config
.connect(crate::usage::tls_connector())
.await
.map_err(|error| classify(&error))?;
let driver = tokio::spawn(async move {
let _ = connection.await;
});
let queried = client.simple_query("SELECT 1").await;
drop(client);
driver.abort();
queried.map(|_| ()).map_err(|error| classify(&error))
}
}
fn classify(error: &tokio_postgres::Error) -> HealthFailure {
use tokio_postgres::error::SqlState;
let category = match error.code() {
Some(code)
if *code == SqlState::INVALID_PASSWORD
|| *code == SqlState::INVALID_AUTHORIZATION_SPECIFICATION
|| *code == SqlState::INSUFFICIENT_PRIVILEGE =>
{
FailureCategory::Denied
}
_ => FailureCategory::Unavailable,
};
HealthFailure::new(category, error.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn an_unreachable_postgres_is_unavailable_and_names_no_dsn() {
let config: Config = "host=127.0.0.1 port=1 user=axond connect_timeout=1"
.parse()
.expect("parsable DSN");
let health = PostgresHealth::new("postgres", config, Duration::from_secs(2));
let failure = health.check().await.expect_err("port 1 refuses");
assert_eq!(failure.category(), FailureCategory::Unavailable);
assert!(!failure.detail().contains("user=axond"), "{failure:?}");
}
}