sunbeam-g2v 0.2.0

//! Health check module for Sunbeam services.
//!
//! Provides liveness and readiness health checks, plus an [`HealthRouter`]
//! that exposes them as axum routes:
//!
//! - `GET /health/live`  — always 200; process is up.
//! - `GET /health/ready` — 200 if all registered checks pass, 503 otherwise.
//!
//! # Example
//!
//! ```rust,no_run
//! # #[cfg(all(feature = "axum", feature = "sqlx-postgres"))]
//! # async fn example(pool: sqlx::PgPool) {
//! use std::sync::Arc;
//! use sunbeam_g2v::health::{DatabaseHealthCheck, HealthRouter};
//!
//! let router = HealthRouter::new()
//!     .with_check(Arc::new(DatabaseHealthCheck::new(pool)))
//!     .into_axum_router();
//! # }
//! ```

pub mod router;

use crate::error::ServiceResult;
use std::sync::Arc;

pub use router::HealthRouter;

// ============================================================================
// Result type
// ============================================================================

/// Result of a single health check.
#[derive(Debug, Clone)]
pub struct HealthCheckResult {
    /// Whether the component is healthy.
    pub healthy: bool,
    /// Optional human-readable details (required when unhealthy).
    pub details: Option<String>,
}

impl HealthCheckResult {
    /// Create a healthy result.
    pub fn healthy() -> Self {
        Self {
            healthy: true,
            details: None,
        }
    }

    /// Create an unhealthy result with a reason.
    pub fn unhealthy(details: impl Into<String>) -> Self {
        Self {
            healthy: false,
            details: Some(details.into()),
        }
    }

    /// Returns `true` when healthy.
    pub fn is_healthy(&self) -> bool {
        self.healthy
    }
}

// ============================================================================
// Trait
// ============================================================================

/// A single health check that can be polled asynchronously.
///
/// Implementations must be `Send + Sync + 'static` so they can be stored
/// behind `Arc<dyn HealthCheck>` and awaited from any task.
pub trait HealthCheck: Send + Sync + 'static {
    /// Short name used in the readiness JSON response (e.g. `"database"`).
    fn name(&self) -> &str;

    /// Run the check and return a result.
    fn check(
        &self,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static>,
    >;
}

// ============================================================================
// DatabaseHealthCheck
// ============================================================================

/// Health check that issues `SELECT 1` against a Postgres pool.
///
/// The query is wrapped in a 2-second timeout. Any error (connection failure,
/// timeout, etc.) is returned as an unhealthy result rather than propagating
/// as a `ServiceError`, so the readiness endpoint always returns a structured
/// JSON payload.
#[cfg(feature = "sqlx-postgres")]
pub struct DatabaseHealthCheck {
    pool: sqlx::PgPool,
}

#[cfg(feature = "sqlx-postgres")]
impl DatabaseHealthCheck {
    /// Create a check from a live pool.
    pub fn new(pool: sqlx::PgPool) -> Self {
        Self { pool }
    }
}

#[cfg(feature = "sqlx-postgres")]
impl HealthCheck for DatabaseHealthCheck {
    fn name(&self) -> &str {
        "database"
    }

    fn check(
        &self,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static>,
    > {
        let pool = self.pool.clone();
        Box::pin(async move {
            let timeout = tokio::time::timeout(
                std::time::Duration::from_secs(2),
                sqlx::query("SELECT 1").execute(&pool),
            )
            .await;

            match timeout {
                Ok(Ok(_)) => Ok(HealthCheckResult::healthy()),
                Ok(Err(e)) => Ok(HealthCheckResult::unhealthy(format!("query failed: {e}"))),
                Err(_) => Ok(HealthCheckResult::unhealthy("timed out after 2s")),
            }
        })
    }
}

// ============================================================================
// KetoHealthCheck
// ============================================================================

/// Health check for the Ory Keto gRPC service.
///
/// Issues a `CheckService.Check` against a sentinel namespace/object. Any
/// successful gRPC response (including "denied") means Keto is reachable and
/// healthy. A transport error means Keto is unreachable.
#[cfg(feature = "keto")]
pub struct KetoHealthCheck {
    client: crate::middleware::auth::KetoClient,
}

#[cfg(feature = "keto")]
impl KetoHealthCheck {
    /// Create a check from an existing [`KetoClient`][crate::middleware::auth::KetoClient].
    pub fn new(client: crate::middleware::auth::KetoClient) -> Self {
        Self { client }
    }
}

#[cfg(feature = "keto")]
impl HealthCheck for KetoHealthCheck {
    fn name(&self) -> &str {
        "keto"
    }

    fn check(
        &self,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static>,
    > {
        let client = self.client.clone();
        Box::pin(async move {
            // A successful response (allowed OR denied) proves Keto is up.
            // Only a transport/network error is treated as unhealthy.
            let timeout = tokio::time::timeout(
                std::time::Duration::from_secs(2),
                client.check_permission("__health__", "__probe__", "probe", "__health_probe__"),
            )
            .await;

            match timeout {
                // `Ok(true)` or `Ok(false)` — Keto responded, it's healthy.
                Ok(Ok(_)) => Ok(HealthCheckResult::healthy()),
                // ServiceError::Internal means Keto returned a gRPC error other
                // than PermissionDenied, which still proves it's reachable.
                // Re-map to healthy because the probe namespace won't exist.
                Ok(Err(crate::error::ServiceError::Internal(_))) => {
                    Ok(HealthCheckResult::healthy())
                }
                Ok(Err(e)) => Ok(HealthCheckResult::unhealthy(format!("keto error: {e}"))),
                Err(_) => Ok(HealthCheckResult::unhealthy("timed out after 2s")),
            }
        })
    }
}

// ============================================================================
// NatsHealthCheck
// ============================================================================

/// Health check for a NATS connection.
///
/// Returns healthy if `client.connection_state()` is `Connected`.
#[cfg(feature = "nats")]
pub struct NatsHealthCheck {
    client: async_nats::Client,
}

#[cfg(feature = "nats")]
impl NatsHealthCheck {
    /// Create a check from an existing NATS client.
    pub fn new(client: async_nats::Client) -> Self {
        Self { client }
    }
}

#[cfg(feature = "nats")]
impl HealthCheck for NatsHealthCheck {
    fn name(&self) -> &str {
        "nats"
    }

    fn check(
        &self,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static>,
    > {
        use async_nats::connection::State;
        let state = self.client.connection_state();
        Box::pin(async move {
            if state == State::Connected {
                Ok(HealthCheckResult::healthy())
            } else {
                Ok(HealthCheckResult::unhealthy(format!(
                    "NATS connection state: {state:?}"
                )))
            }
        })
    }
}

// ============================================================================
// HttpHealthCheck
// ============================================================================

/// Health check that GETs a URL and asserts a 2xx response.
///
/// Uses a 2-second timeout. Useful for poking upstream HTTP services.
pub struct HttpHealthCheck {
    client: reqwest::Client,
    url: String,
    name: String,
}

impl HttpHealthCheck {
    /// Create a check with an explicit name, client, and URL.
    pub fn new(name: impl Into<String>, client: reqwest::Client, url: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            client,
            url: url.into(),
        }
    }

    /// Convenience constructor that builds a default reqwest client.
    pub fn with_url(name: impl Into<String>, url: impl Into<String>) -> Self {
        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(2))
            .build()
            .expect("build reqwest client for HttpHealthCheck");
        Self::new(name, client, url)
    }
}

impl HealthCheck for HttpHealthCheck {
    fn name(&self) -> &str {
        &self.name
    }

    fn check(
        &self,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static>,
    > {
        let client = self.client.clone();
        let url = self.url.clone();
        Box::pin(async move {
            let timeout = tokio::time::timeout(
                std::time::Duration::from_secs(2),
                client.get(&url).send(),
            )
            .await;

            match timeout {
                Ok(Ok(resp)) if resp.status().is_success() => Ok(HealthCheckResult::healthy()),
                Ok(Ok(resp)) => Ok(HealthCheckResult::unhealthy(format!(
                    "HTTP {} from {url}",
                    resp.status()
                ))),
                Ok(Err(e)) => Ok(HealthCheckResult::unhealthy(format!("request error: {e}"))),
                Err(_) => Ok(HealthCheckResult::unhealthy("timed out after 2s")),
            }
        })
    }
}

// ============================================================================
// CompositeHealthCheck
// ============================================================================

/// Runs a sequence of checks and short-circuits on the first failure.
///
/// All checks are run in order; the first unhealthy result is returned
/// immediately. If all checks pass, a healthy result is returned.
pub struct CompositeHealthCheck {
    checks: Vec<Arc<dyn HealthCheck>>,
}

impl CompositeHealthCheck {
    /// Create an empty composite check.
    pub fn new() -> Self {
        Self { checks: Vec::new() }
    }

    /// Add a check to the sequence.
    pub fn add_check(mut self, check: Arc<dyn HealthCheck>) -> Self {
        self.checks.push(check);
        self
    }
}

impl Default for CompositeHealthCheck {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for CompositeHealthCheck {
    fn clone(&self) -> Self {
        Self {
            checks: self.checks.clone(),
        }
    }
}

impl HealthCheck for CompositeHealthCheck {
    fn name(&self) -> &str {
        "composite"
    }

    fn check(
        &self,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static>,
    > {
        let checks = self.checks.clone();
        Box::pin(async move {
            for check in &checks {
                let result = check.check().await?;
                if !result.is_healthy() {
                    return Ok(result);
                }
            }
            Ok(HealthCheckResult::healthy())
        })
    }
}

// ============================================================================
// Unit tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;

    // -----------------------------------------------------------------------
    // HealthCheckResult
    // -----------------------------------------------------------------------

    #[test]
    fn test_health_check_result_healthy() {
        let result = HealthCheckResult::healthy();
        assert!(result.is_healthy());
        assert!(result.details.is_none());
    }

    #[test]
    fn test_health_check_result_unhealthy() {
        let result = HealthCheckResult::unhealthy("Connection failed");
        assert!(!result.is_healthy());
        assert_eq!(result.details, Some("Connection failed".to_string()));
    }

    // -----------------------------------------------------------------------
    // CompositeHealthCheck
    // -----------------------------------------------------------------------

    struct AlwaysHealthy;
    impl HealthCheck for AlwaysHealthy {
        fn name(&self) -> &str {
            "always-healthy"
        }
        fn check(
            &self,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static,
            >,
        > {
            Box::pin(async { Ok(HealthCheckResult::healthy()) })
        }
    }

    struct AlwaysUnhealthy;
    impl HealthCheck for AlwaysUnhealthy {
        fn name(&self) -> &str {
            "always-unhealthy"
        }
        fn check(
            &self,
        ) -> std::pin::Pin<
            Box<
                dyn std::future::Future<Output = ServiceResult<HealthCheckResult>> + Send + 'static,
            >,
        > {
            Box::pin(async { Ok(HealthCheckResult::unhealthy("injected failure")) })
        }
    }

    #[tokio::test]
    async fn test_composite_all_healthy() {
        let composite = CompositeHealthCheck::new()
            .add_check(Arc::new(AlwaysHealthy))
            .add_check(Arc::new(AlwaysHealthy));
        let result = composite.check().await.unwrap();
        assert!(result.is_healthy());
    }

    #[tokio::test]
    async fn test_composite_short_circuits_on_first_failure() {
        // First check fails; second check is healthy but should never be reached.
        let composite = CompositeHealthCheck::new()
            .add_check(Arc::new(AlwaysUnhealthy))
            .add_check(Arc::new(AlwaysHealthy));
        let result = composite.check().await.unwrap();
        assert!(!result.is_healthy());
        assert_eq!(result.details, Some("injected failure".to_string()));
    }

    #[tokio::test]
    async fn test_composite_empty_is_healthy() {
        let composite = CompositeHealthCheck::new();
        let result = composite.check().await.unwrap();
        assert!(result.is_healthy());
    }

    // -----------------------------------------------------------------------
    // HealthRouter builder
    // -----------------------------------------------------------------------

    #[test]
    fn test_health_router_new_is_empty() {
        let r = HealthRouter::new();
        assert!(r.checks.is_empty());
    }

    #[test]
    fn test_health_router_with_check_adds() {
        let r = HealthRouter::new()
            .with_check(Arc::new(AlwaysHealthy))
            .with_check(Arc::new(AlwaysUnhealthy));
        assert_eq!(r.checks.len(), 2);
    }

    // -----------------------------------------------------------------------
    // Route handler behaviour
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn test_live_always_200() {
        use crate::testing::TestHarness;

        let health_router = HealthRouter::new()
            .with_check(Arc::new(AlwaysUnhealthy)) // checks shouldn't matter for /live
            .into_axum_router();

        let mut harness = TestHarness::new().with_routes(health_router);
        harness.start().await.unwrap();

        let resp = harness.get("/health/live").await.unwrap();
        assert_eq!(resp.status, 200, "live must always be 200");
        harness.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_ready_200_when_all_pass() {
        use crate::testing::TestHarness;

        let health_router = HealthRouter::new()
            .with_check(Arc::new(AlwaysHealthy))
            .with_check(Arc::new(AlwaysHealthy))
            .into_axum_router();

        let mut harness = TestHarness::new().with_routes(health_router);
        harness.start().await.unwrap();

        let resp = harness.get("/health/ready").await.unwrap();
        assert_eq!(resp.status, 200);

        let body: serde_json::Value = serde_json::from_slice(&resp.body).unwrap();
        assert_eq!(body["status"], "ok");
        assert!(body["checks"].is_array());
        harness.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_ready_503_when_one_fails() {
        use crate::testing::TestHarness;

        let health_router = HealthRouter::new()
            .with_check(Arc::new(AlwaysHealthy))
            .with_check(Arc::new(AlwaysUnhealthy))
            .into_axum_router();

        let mut harness = TestHarness::new().with_routes(health_router);
        harness.start().await.unwrap();

        let resp = harness.get("/health/ready").await.unwrap();
        assert_eq!(resp.status, 503);

        let body: serde_json::Value = serde_json::from_slice(&resp.body).unwrap();
        assert_eq!(body["status"], "degraded");
        harness.stop().await.unwrap();
    }
}