fix(store): Connection opts for specific databases (#384)

FuelLabs · Jan 27, 2025 · e2560d4 · e2560d4
1 parent 5973927
commit e2560d4
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 10 deletions.
diff --git a/.env.sample b/.env.sample
@@ -1,4 +1,5 @@
 # General Configuration
+DB_TYPE=Cockroach
 DB_POOL_SIZE=100
 DB_TIMEOUT=240
 STORE_PAGINATION_LIMIT=100

diff --git a/cluster/charts/fuel-streams/Chart.yaml b/cluster/charts/fuel-streams/Chart.yaml
@@ -2,7 +2,7 @@ apiVersion: v2
 appVersion: "1.0"
 description: A Helm chart for Kubernetes
 name: fuel-streams
-version: 0.9.3
+version: 0.9.4
 dependencies:
   - name: nats
     version: 1.2.9

diff --git a/cluster/charts/fuel-streams/values.yaml b/cluster/charts/fuel-streams/values.yaml
@@ -80,6 +80,7 @@ commonConfigMap:
     NATS_ADMIN_USER: "admin"
     NATS_ADMIN_PASS: "admin"
     # For local purposes only, for production use fuel-streams-keys secret
+    DB_TYPE: "Aurora"
     DATABASE_URL: "postgresql://root:root@fuel-streams-cockroachdb:26257/fuel_streams?sslmode=disable"
 
 # This is a secret that is used for local development

diff --git a/crates/fuel-streams-store/src/db/db_impl.rs b/crates/fuel-streams-store/src/db/db_impl.rs
@@ -60,9 +60,9 @@ impl Default for DbConnectionOpts {
             pool_size: Some(*DB_POOL_SIZE as u32),
             connection_str: dotenvy::var("DATABASE_URL")
                 .expect("DATABASE_URL not set"),
-            statement_timeout: Some(Duration::from_secs(30)),
+            statement_timeout: Some(Duration::from_secs(60)),
             acquire_timeout: Some(Duration::from_secs(10)),
-            idle_timeout: Some(Duration::from_secs(180)),
+            idle_timeout: Some(Duration::from_secs(240)),
             min_connections: Some((*DB_POOL_SIZE as u32) / 4),
         }
     }
@@ -92,19 +92,13 @@ impl Db {
     }
 
     async fn create_pool(opts: &DbConnectionOpts) -> DbResult<Pool<Postgres>> {
-        let statement_timeout =
-            opts.statement_timeout.unwrap_or_default().as_millis();
-        let statement_timeout = &format!("{}", statement_timeout);
         let connections_opts =
             sqlx::postgres::PgConnectOptions::from_str(&opts.connection_str)
                 .map_err(|e| {
                     DbError::Open(sqlx::Error::Configuration(Box::new(e)))
                 })?
                 .application_name("fuel-streams")
-                .options([
-                    ("retry_connect_backoff", "2"),
-                    ("statement_timeout", statement_timeout),
-                ]);
+                .options(Self::connect_opts(&opts));
 
         sqlx::postgres::PgPoolOptions::new()
             .max_connections(opts.pool_size.unwrap_or_default())
@@ -116,4 +110,27 @@ impl Db {
             .await
             .map_err(DbError::Open)
     }
+
+    fn connect_opts(opts: &DbConnectionOpts) -> Vec<(String, String)> {
+        let db_type =
+            dotenvy::var("DB_TYPE").expect("Database type need to be defined");
+        let statement_timeout =
+            opts.statement_timeout.unwrap_or_default().as_millis();
+        let statement_timeout = format!("{}", statement_timeout);
+        let idle_timeout = opts.idle_timeout.unwrap_or_default().as_millis();
+        let idle_timeout = format!("{}", idle_timeout);
+        let mut connect_opts = vec![];
+        if db_type == "Aurora" {
+            connect_opts.push((
+                "idle_in_transaction_session_timeout".to_string(),
+                idle_timeout,
+            ));
+        }
+        if db_type == "Cockroach" {
+            connect_opts
+                .push(("retry_connect_backoff".to_string(), "2".to_string()));
+        }
+        connect_opts.push(("statement_timeout".to_string(), statement_timeout));
+        connect_opts
+    }
 }