From 356d91de7ad27b0ecfaabc1df9d7f639edfa52d7 Mon Sep 17 00:00:00 2001
From: Dimitri <dimitridumonet@gmail.com>
Date: Tue, 17 Sep 2024 15:41:17 +0700
Subject: [PATCH] Force db initialization on startup

---
 Cargo.lock                                    |  1 +
 packages/relayer/src/database.rs              | 25 ++++++++++++++-
 packages/relayer/src/lib.rs                   | 32 ++++++++++++-------
 .../relayer/src/modules/web_server/server.rs  | 17 ++++++++++
 4 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ac64a597..212f8687 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4141,6 +4141,7 @@ dependencies = [
 [[package]]
 name = "relayer-utils"
 version = "0.3.7"
+source = "git+https://github.com/zkemail/relayer-utils.git?rev=94d78d6#94d78d67862b6d6c15bebac66d184c7557f6aff5"
 dependencies = [
  "anyhow",
  "base64 0.21.7",
diff --git a/packages/relayer/src/database.rs b/packages/relayer/src/database.rs
index 5dd0fbb9..1af6a220 100644
--- a/packages/relayer/src/database.rs
+++ b/packages/relayer/src/database.rs
@@ -73,6 +73,28 @@ impl Database {
         Ok(())
     }
 
+    pub(crate) async fn test_db_connection(&self) -> Result<()> {
+        // Try up to 3 times
+        for i in 1..4 {
+            match sqlx::query("SELECT 1").execute(&self.db).await {
+                Ok(_) => {
+                    info!(LOG, "Connected successfully to database");
+                    return Ok(());
+                }
+                Err(e) => {
+                    error!(
+                        LOG,
+                        "Failed to initialize connection to the database: {:?}. Retrying...", e
+                    );
+                    tokio::time::sleep(Duration::from_secs(i * i)).await;
+                }
+            }
+        }
+        Err(anyhow::anyhow!(
+            "Failed to initialize database connection after 3 attempts"
+        ))
+    }
+
     pub(crate) async fn get_credentials(&self, account_code: &str) -> Result<Option<Credentials>> {
         let row = sqlx::query("SELECT * FROM credentials WHERE account_code = $1")
             .bind(account_code)
@@ -326,6 +348,7 @@ impl Database {
         &self,
         row: &Request,
     ) -> std::result::Result<(), DatabaseError> {
+        let request_id = row.request_id;
         let row = sqlx::query(
             "INSERT INTO requests (request_id, account_eth_addr, controller_eth_addr, guardian_email_addr, is_for_recovery, template_idx, is_processed, is_success, email_nullifier, account_salt) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING *",
         )
@@ -342,7 +365,7 @@ impl Database {
         .fetch_one(&self.db)
         .await
         .map_err(|e| DatabaseError::new("Failed to insert request", e))?;
-        info!(LOG, "Request inserted");
+        info!(LOG, "Request inserted with request_id: {}", request_id);
         Ok(())
     }
 }
diff --git a/packages/relayer/src/lib.rs b/packages/relayer/src/lib.rs
index 5cc01621..42348138 100644
--- a/packages/relayer/src/lib.rs
+++ b/packages/relayer/src/lib.rs
@@ -19,7 +19,7 @@ pub use modules::*;
 use relayer_utils::LOG;
 pub use strings::*;
 
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, OnceCell};
 
 use anyhow::{anyhow, Result};
 use dotenv::dotenv;
@@ -44,17 +44,27 @@ pub static EMAIL_TEMPLATES: OnceLock<String> = OnceLock::new();
 pub static RELAYER_EMAIL_ADDRESS: OnceLock<String> = OnceLock::new();
 pub static SMTP_SERVER: OnceLock<String> = OnceLock::new();
 
+static DB_CELL: OnceCell<Arc<Database>> = OnceCell::const_new();
+
+struct DBWrapper;
+
+impl DBWrapper {
+    fn get() -> &'static Arc<Database> {
+        DB_CELL.get().expect("Database not initialized")
+    }
+}
+
+impl std::ops::Deref for DBWrapper {
+    type Target = Database;
+
+    fn deref(&self) -> &Self::Target {
+        &**Self::get()
+    }
+}
+
+static DB: DBWrapper = DBWrapper;
+
 lazy_static! {
-    pub static ref DB: Arc<Database> = {
-        dotenv().ok();
-        let db = tokio::task::block_in_place(|| {
-            tokio::runtime::Runtime::new()
-                .unwrap()
-                .block_on(Database::open(&env::var(DATABASE_PATH_KEY).unwrap()))
-        })
-        .unwrap();
-        Arc::new(db)
-    };
     pub static ref CLIENT: Arc<ChainClient> = {
         dotenv().ok();
         let client = tokio::task::block_in_place(|| {
diff --git a/packages/relayer/src/modules/web_server/server.rs b/packages/relayer/src/modules/web_server/server.rs
index fbf85e77..f96afb12 100644
--- a/packages/relayer/src/modules/web_server/server.rs
+++ b/packages/relayer/src/modules/web_server/server.rs
@@ -6,6 +6,23 @@ use tower_http::cors::{AllowHeaders, AllowMethods, Any, CorsLayer};
 pub async fn run_server() -> Result<()> {
     let addr = WEB_SERVER_ADDRESS.get().unwrap();
 
+    DB_CELL
+        .get_or_init(|| async {
+            dotenv::dotenv().ok();
+            let db = Database::open(&std::env::var("DATABASE_URL").unwrap())
+                .await
+                .unwrap();
+            Arc::new(db)
+        })
+        .await;
+
+    info!(LOG, "Testing connection to database");
+    if let Err(e) = DB.test_db_connection().await {
+        error!(LOG, "Failed to initialize db with e: {}", e);
+        panic!("Forcing panic, since connection to DB could not be established");
+    };
+    info!(LOG, "Testing connection to database successfull");
+
     let mut app = Router::new()
         .route(
             "/api/echo",