feat: use db orchestrator in gateway + remove mutable reference to self

Author: MarcosNicolau

aggregation_mode/Cargo.lock

@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+serde =  { workspace = true }
 tokio = { version = "1"}
 sqlx = { version = "0.8", features = [ "runtime-tokio", "postgres", "migrate" ] }
 backon = "1.2.0"

aggregation_mode/db/Cargo.toml

@@ -93,27 +93,23 @@ impl DbOrchestartor {
         })
     }
 
-    pub async fn write<T, Q, Fut>(&mut self, query: Q) -> Result<T, sqlx::Error>
+    pub async fn write<T, Q, Fut>(&self, query: Q) -> Result<T, sqlx::Error>
     where
         Q: Fn(Pool<Postgres>) -> Fut,
         Fut: Future<Output = Result<T, sqlx::Error>>,
     {
         self.query::<T, Q, Fut>(query, Operation::Write).await
     }
 
-    pub async fn read<T, Q, Fut>(&mut self, query: Q) -> Result<T, sqlx::Error>
+    pub async fn read<T, Q, Fut>(&self, query: Q) -> Result<T, sqlx::Error>
     where
         Q: Fn(Pool<Postgres>) -> Fut,
         Fut: Future<Output = Result<T, sqlx::Error>>,
     {
         self.query::<T, Q, Fut>(query, Operation::Read).await
     }
 
-    async fn query<T, Q, Fut>(
-        &mut self,
-        query_fn: Q,
-        operation: Operation,
-    ) -> Result<T, sqlx::Error>
+    async fn query<T, Q, Fut>(&self, query_fn: Q, operation: Operation) -> Result<T, sqlx::Error>
     where
         Q: Fn(Pool<Postgres>) -> Fut,
         Fut: Future<Output = Result<T, sqlx::Error>>,
@@ -151,7 +147,7 @@ impl DbOrchestartor {
     }
 
     async fn execute_once<T, Q, Fut>(
-        &mut self,
+        &self,
         query_fn: &Q,
         operation: Operation,
     ) -> Result<T, RetryError<sqlx::Error>>

aggregation_mode/db/src/orchestrator.rs

@@ -4,7 +4,7 @@ use sqlx::{
     Type,
 };
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Type)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Type, serde::Serialize)]
 #[sqlx(type_name = "task_status", rename_all = "lowercase")]
 pub enum TaskStatus {
     Pending,

aggregation_mode/db/src/types.rs

@@ -10,6 +10,7 @@ serde_yaml = { workspace = true }
 agg_mode_sdk = { path = "../sdk"}
 aligned-sdk = { workspace = true }
 sp1-sdk = { workspace = true }
+db = { workspace = true }
 tracing = { version = "0.1", features = ["log"] }
 tracing-subscriber = { version = "0.3.0", features = ["env-filter"] }
 bincode = "1.3.3"

aggregation_mode/gateway/Cargo.toml

@@ -1,111 +1,121 @@
-use sqlx::{
-    postgres::PgPoolOptions,
-    types::{BigDecimal, Uuid},
-    Pool, Postgres,
-};
+use db::{orchestrator::DbOrchestartor, retry::RetryConfig};
+use sqlx::types::{BigDecimal, Uuid};
+
+use crate::types::Receipt;
 
 #[derive(Clone, Debug)]
 pub struct Db {
-    pool: Pool<Postgres>,
+    orchestrator: DbOrchestartor,
 }
 
 #[derive(Debug, Clone)]
 pub enum DbError {
     ConnectError(String),
 }
 
-#[derive(Debug, Clone, sqlx::Type, serde::Serialize)]
-#[sqlx(type_name = "task_status")]
-#[sqlx(rename_all = "lowercase")]
-pub enum TaskStatus {
-    Pending,
-    Processing,
-    Verified,
-}
-
-#[derive(Debug, Clone, sqlx::FromRow, sqlx::Type, serde::Serialize)]
-pub struct Receipt {
-    pub status: TaskStatus,
-    pub merkle_path: Option<Vec<u8>>,
-    pub nonce: i64,
-    pub address: String,
-}
-
 impl Db {
-    pub async fn try_new(connection_url: &str) -> Result<Self, DbError> {
-        let pool = PgPoolOptions::new()
-            .max_connections(5)
-            .connect(connection_url)
-            .await
-            .map_err(|e| DbError::ConnectError(e.to_string()))?;
+    pub async fn try_new(connection_urls: &[&str]) -> Result<Self, DbError> {
+        let orchestrator = DbOrchestartor::try_new(
+            connection_urls,
+            RetryConfig {
+                factor: 0.0,
+                max_delay_seconds: 0,
+                max_times: 0,
+                min_delay_millis: 0,
+            },
+        )
+        .map_err(|e| DbError::ConnectError(e.to_string()))?;
 
-        Ok(Self { pool })
+        Ok(Self { orchestrator })
     }
 
     pub async fn count_tasks_by_address(&self, address: &str) -> Result<i64, sqlx::Error> {
-        let (count,) = sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM tasks WHERE address = $1")
-            .bind(address.to_lowercase())
-            .fetch_one(&self.pool)
-            .await?;
+        self.orchestrator
+            .read(async |pool| {
+                let (count,) =
+                    sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM tasks WHERE address = $1")
+                        .bind(address.to_lowercase())
+                        .fetch_one(&pool)
+                        .await?;
 
-        Ok(count)
+                Ok(count)
+            })
+            .await
     }
 
     pub async fn get_merkle_path_by_task_id(
         &self,
         task_id: Uuid,
     ) -> Result<Option<Vec<u8>>, sqlx::Error> {
-        sqlx::query_scalar::<_, Option<Vec<u8>>>("SELECT merkle_path FROM tasks WHERE task_id = $1")
-            .bind(task_id)
-            .fetch_optional(&self.pool)
+        self.orchestrator
+            .read(async |pool| {
+                sqlx::query_scalar::<_, Option<Vec<u8>>>(
+                    "SELECT merkle_path FROM tasks WHERE task_id = $1",
+                )
+                .bind(task_id)
+                .fetch_optional(&pool)
+                .await
+                .map(|res| res.flatten())
+            })
             .await
-            .map(|res| res.flatten())
     }
 
     pub async fn get_tasks_by_address_and_nonce(
         &self,
         address: &str,
         nonce: i64,
     ) -> Result<Vec<Receipt>, sqlx::Error> {
-        sqlx::query_as::<_, Receipt>(
-            "SELECT status,merkle_path,nonce,address FROM tasks
-                WHERE address = $1
-                AND nonce = $2
-                ORDER BY nonce DESC",
-        )
-        .bind(address.to_lowercase())
-        .bind(nonce)
-        .fetch_all(&self.pool)
-        .await
+        self.orchestrator
+            .read(async |pool| {
+                sqlx::query_as::<_, Receipt>(
+                    "SELECT status,merkle_path,nonce,address FROM tasks
+                    WHERE address = $1
+                    AND nonce = $2
+                    ORDER BY nonce DESC",
+                )
+                .bind(address.to_lowercase())
+                .bind(nonce)
+                .fetch_all(&pool)
+                .await
+            })
+            .await
     }
 
     pub async fn get_tasks_by_address_with_limit(
         &self,
         address: &str,
         limit: i64,
     ) -> Result<Vec<Receipt>, sqlx::Error> {
-        sqlx::query_as::<_, Receipt>(
-            "SELECT status,merkle_path,nonce,address FROM tasks
-                WHERE address = $1
-                ORDER BY nonce DESC
-                LIMIT $2",
-        )
-        .bind(address.to_lowercase())
-        .bind(limit)
-        .fetch_all(&self.pool)
-        .await
+        self.orchestrator
+            .read(async |pool| {
+                sqlx::query_as::<_, Receipt>(
+                    "SELECT status,merkle_path,nonce,address FROM tasks
+                    WHERE address = $1
+                    ORDER BY nonce DESC
+                    LIMIT $2",
+                )
+                .bind(address.to_lowercase())
+                .bind(limit)
+                .fetch_all(&pool)
+                .await
+            })
+            .await
     }
 
     pub async fn get_daily_tasks_by_address(&self, address: &str) -> Result<i64, sqlx::Error> {
-        sqlx::query_scalar::<_, i64>(
-            "SELECT COUNT(*)
-            FROM tasks
-            WHERE address = $1
-            AND inserted_at::date = CURRENT_DATE",
-        )
-        .bind(address.to_lowercase())
-        .fetch_one(&self.pool)
-        .await
+        self.orchestrator
+            .read(async |pool| {
+                sqlx::query_scalar::<_, i64>(
+                    "SELECT COUNT(*)
+                FROM tasks
+                WHERE address = $1
+                AND inserted_at::date = CURRENT_DATE",
+                )
+                .bind(address.to_lowercase())
+                .fetch_one(&pool)
+                .await
+            })
+            .await
     }
 
     pub async fn insert_task(
@@ -117,41 +127,49 @@ impl Db {
         merkle_path: Option<&[u8]>,
         nonce: i64,
     ) -> Result<Uuid, sqlx::Error> {
-        sqlx::query_scalar::<_, Uuid>(
-            "INSERT INTO tasks (
-                address,
-                proving_system_id,
-                proof,
-                program_commitment,
-                merkle_path,
-                nonce
-            ) VALUES ($1, $2, $3, $4, $5, $6)
-            RETURNING task_id",
-        )
-        .bind(address.to_lowercase())
-        .bind(proving_system_id)
-        .bind(proof)
-        .bind(program_commitment)
-        .bind(merkle_path)
-        .bind(nonce)
-        .fetch_one(&self.pool)
-        .await
+        self.orchestrator
+            .write(async |pool| {
+                sqlx::query_scalar::<_, Uuid>(
+                    "INSERT INTO tasks (
+                        address,
+                        proving_system_id,
+                        proof,
+                        program_commitment,
+                        merkle_path,
+                        nonce
+                    ) VALUES ($1, $2, $3, $4, $5, $6)
+                    RETURNING task_id",
+                )
+                .bind(address.to_lowercase())
+                .bind(proving_system_id)
+                .bind(proof)
+                .bind(program_commitment)
+                .bind(merkle_path)
+                .bind(nonce)
+                .fetch_one(&pool)
+                .await
+            })
+            .await
     }
 
     pub async fn has_active_payment_event(
         &self,
         address: &str,
         epoch: BigDecimal,
     ) -> Result<bool, sqlx::Error> {
-        sqlx::query_scalar::<_, bool>(
-            "SELECT EXISTS (
-                SELECT 1 FROM payment_events
-                WHERE address = $1 AND started_at < $2 AND $2 < valid_until
-            )",
-        )
-        .bind(address.to_lowercase())
-        .bind(epoch)
-        .fetch_one(&self.pool)
-        .await
+        self.orchestrator
+            .read(async |pool| {
+                sqlx::query_scalar::<_, bool>(
+                    "SELECT EXISTS (
+                    SELECT 1 FROM payment_events
+                    WHERE address = $1 AND started_at < $2 AND $2 < valid_until
+                )",
+                )
+                .bind(address.to_lowercase())
+                .bind(&epoch)
+                .fetch_one(&pool)
+                .await
+            })
+            .await
     }
 }

aggregation_mode/gateway/src/db.rs

@@ -84,6 +84,7 @@ impl GatewayServer {
                 .json(AppResponse::new_unsucessfull("Internal server error", 500));
         };
 
+        // TODO: how to fix the mutable thing
         let state = state.get_ref();
         match state.db.count_tasks_by_address(&address).await {
             Ok(count) => HttpResponse::Ok().json(AppResponse::new_sucessfull(serde_json::json!(

aggregation_mode/gateway/src/http.rs

@@ -1,9 +1,8 @@
 use actix_multipart::form::{tempfile::TempFile, text::Text, MultipartForm};
+use db::types::TaskStatus;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
-use crate::db::TaskStatus;
-
 #[derive(Serialize, Deserialize)]
 pub(super) struct AppResponse {
     status: u16,
@@ -62,3 +61,11 @@ pub struct GetReceiptsResponse {
     pub nonce: i64,
     pub address: String,
 }
+
+#[derive(Debug, Clone, sqlx::FromRow, sqlx::Type, serde::Serialize)]
+pub struct Receipt {
+    pub status: TaskStatus,
+    pub merkle_path: Option<Vec<u8>>,
+    pub nonce: i64,
+    pub address: String,
+}