From 0f17d1cbae25b3f0908de89a1a4a4c1fdf3912f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Charlotte=20=F0=9F=A6=9D=20Delenk?= <lotte@chir.rs>
Date: Thu, 28 Nov 2024 15:44:45 +0100
Subject: [PATCH] actually make readyz useful

---
 ...e30de43484ab9c574f6ae0c01331419bf9c31.json | 20 +++++++++++++
 chir-rs-db/src/lib.rs                         | 29 +++++++++++++++++--
 chir-rs-http/src/lib.rs                       | 20 +++++++++++--
 3 files changed, 63 insertions(+), 6 deletions(-)
 create mode 100644 chir-rs-db/.sqlx/query-882da6c909603e0916ea9971dafe30de43484ab9c574f6ae0c01331419bf9c31.json

diff --git a/chir-rs-db/.sqlx/query-882da6c909603e0916ea9971dafe30de43484ab9c574f6ae0c01331419bf9c31.json b/chir-rs-db/.sqlx/query-882da6c909603e0916ea9971dafe30de43484ab9c574f6ae0c01331419bf9c31.json
new file mode 100644
index 0000000..dca2dac
--- /dev/null
+++ b/chir-rs-db/.sqlx/query-882da6c909603e0916ea9971dafe30de43484ab9c574f6ae0c01331419bf9c31.json
@@ -0,0 +1,20 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT 1 as running",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "running",
+        "type_info": "Int4"
+      }
+    ],
+    "parameters": {
+      "Left": []
+    },
+    "nullable": [
+      null
+    ]
+  },
+  "hash": "882da6c909603e0916ea9971dafe30de43484ab9c574f6ae0c01331419bf9c31"
+}
diff --git a/chir-rs-db/src/lib.rs b/chir-rs-db/src/lib.rs
index 7d44fef..28f64ab 100644
--- a/chir-rs-db/src/lib.rs
+++ b/chir-rs-db/src/lib.rs
@@ -1,9 +1,9 @@
 //! Chir.rs database models
 
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
-use eyre::{Context, Result};
-use sqlx::{migrate, PgPool};
+use eyre::{eyre, Context, Result};
+use sqlx::{migrate, query, PgPool};
 use tracing::instrument;
 
 pub mod file;
@@ -13,6 +13,29 @@ pub mod file;
 #[repr(transparent)]
 pub struct Database(Arc<PgPool>);
 
+impl Database {
+    /// This function verifies an active connection to the database.
+    ///
+    /// # Errors
+    /// This function returns an error if the database connection has failed or a timeout of 1s occurred
+    #[instrument]
+    pub async fn ping(&self) -> Result<()> {
+        let fut = async {
+            match query!("SELECT 1 as running").fetch_one(&*self.0).await {
+                Ok(v) if v.running == Some(1) => Ok::<_, eyre::Report>(()),
+                Err(e) => Err(e).context("Checking for readiness"),
+                r => Err(eyre!("Unknown database response: {r:#?}")),
+            }
+        };
+
+        tokio::time::timeout(Duration::from_secs(1), fut)
+            .await
+            .context("Awaiting a ping")??;
+
+        Ok(())
+    }
+}
+
 /// Opens the database
 ///
 /// # Errors
diff --git a/chir-rs-http/src/lib.rs b/chir-rs-http/src/lib.rs
index 820ccf0..cdc1ee0 100644
--- a/chir-rs-http/src/lib.rs
+++ b/chir-rs-http/src/lib.rs
@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use axum::{
     extract::{MatchedPath, Request, State},
-    http::Response,
+    http::{Response, StatusCode},
     response::IntoResponse,
     routing::get,
     Router,
@@ -17,7 +17,7 @@ use chir_rs_http_api::{axum::bincode::Bincode, readiness::ReadyState};
 use eyre::{Context, Result};
 use tokio::net::TcpListener;
 use tower_http::trace::TraceLayer;
-use tracing::{info, info_span};
+use tracing::{error, info, info_span};
 
 /// Application state
 #[derive(Clone, Debug)]
@@ -48,7 +48,21 @@ pub async fn main(cfg: Arc<ChirRs>, db: Database, castore: CaStore) -> Result<()
     let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();
     let app = Router::new()
         // Routes here
-        .route("/.api/readyz", get(|| async { Bincode(ReadyState::Ready) }))
+        .route(
+            "/.api/readyz",
+            get(|State(state): State<AppState>| async move {
+                match state.db.ping().await {
+                    Ok(()) => (StatusCode::OK, Bincode(ReadyState::Ready)),
+                    Err(e) => {
+                        error!("Database is not responding: {e:?}");
+                        (
+                            StatusCode::INTERNAL_SERVER_ERROR,
+                            Bincode(ReadyState::NotReady),
+                        )
+                    }
+                }
+            }),
+        )
         .route(
             "/.api/metrics",
             get(|| async move { metric_handle.render() }),