cat_gateway/service/utilities/middleware/
db_check.rs

1//! Middleware to verify that there is connection the databases.
2
3use poem::{http::StatusCode, Endpoint, Middleware, Request, Result};
4use tracing::error;
5
6use crate::service::utilities::health::{event_db_is_live, index_db_is_live};
7
8/// Middleware type that returns a response with 503 status code
9/// if any DB stops responding before returning the wrapped endpoint.
10pub(crate) struct DatabaseConnectionCheck;
11
12impl<E: Endpoint> Middleware<E> for DatabaseConnectionCheck {
13    type Output = DatabaseConnectionImpl<E>;
14
15    fn transform(
16        &self,
17        ep: E,
18    ) -> Self::Output {
19        DatabaseConnectionImpl { ep }
20    }
21}
22
23/// The new endpoint type generated by the `DatabaseConnectionCheck`.
24pub(crate) struct DatabaseConnectionImpl<E> {
25    /// Endpoint wrapped by the middleware.
26    ep: E,
27}
28
29impl<E: Endpoint> Endpoint for DatabaseConnectionImpl<E> {
30    type Output = E::Output;
31
32    async fn call(
33        &self,
34        req: Request,
35    ) -> Result<Self::Output> {
36        let req_path = req.uri().path();
37
38        // TODO: find a better way to filter URI paths
39        let is_health_endpoint = req_path.starts_with("/api/v1/health/");
40        let is_metrics_endpoint = req_path == "/metrics";
41
42        if !(is_health_endpoint || is_metrics_endpoint) {
43            if !event_db_is_live() {
44                error!(endpoint_path = %req_path, "Event DB is not live");
45                return Err(StatusCode::SERVICE_UNAVAILABLE.into());
46            }
47            if !index_db_is_live() {
48                error!(endpoint_path = %req_path, "Index DB is not live");
49                return Err(StatusCode::SERVICE_UNAVAILABLE.into());
50            }
51        }
52        self.ep.call(req).await
53    }
54}