Allow configuring dashboard connection pool settings (#635)

levkk · web-flow · commit b0796cab01ac · 2023-05-18T15:07:45.000-07:00
diff --git a/pgml-dashboard/src/lib.rs b/pgml-dashboard/src/lib.rs
@@ -24,6 +24,13 @@ use guards::Cluster;
 use responses::{BadRequest, ResponseOk};
 use sqlx::Executor;
 
+#[derive(Debug, Default, Clone)]
+pub struct ClustersSettings {
+    pub max_connections: u32,
+    pub idle_timeout: u64,
+    pub min_connections: u32,
+}
+
 /// This struct contains information specific to the cluster being displayed in the dashboard.
 ///
 /// The dashboard is built to manage multiple clusters, but the server itself by design is stateless.
@@ -44,13 +51,18 @@ pub struct Clusters {
 }
 
 impl Clusters {
-    pub fn add(&self, cluster_id: i64, database_url: &str) -> anyhow::Result<PgPool> {
+    pub fn add(
+        &self,
+        cluster_id: i64,
+        database_url: &str,
+        settings: ClustersSettings,
+    ) -> anyhow::Result<PgPool> {
         let mut pools = self.pools.lock();
 
         let pool = PgPoolOptions::new()
-            .max_connections(5)
-            .idle_timeout(std::time::Duration::from_millis(15_000))
-            .min_connections(0)
+            .max_connections(settings.max_connections)
+            .idle_timeout(std::time::Duration::from_millis(settings.idle_timeout))
+            .min_connections(settings.min_connections)
             .after_connect(|conn, _meta| {
                 Box::pin(async move {
                     conn.execute("SET application_name = 'pgml_dashboard';")
diff --git a/pgml-dashboard/src/main.rs b/pgml-dashboard/src/main.rs
@@ -11,8 +11,18 @@ async fn main() {
     dotenv::dotenv().ok();
 
     let clusters = pgml_dashboard::Clusters::new();
+    let settings = pgml_dashboard::ClustersSettings {
+        min_connections: 0,
+        max_connections: 5,
+        idle_timeout: 15_000,
+    };
+
     clusters
-        .add(-1, &pgml_dashboard::guards::default_database_url())
+        .add(
+            -1,
+            &pgml_dashboard::guards::default_database_url(),
+            settings,
+        )
         .unwrap();
 
     pgml_dashboard::migrate(&clusters.get(-1).unwrap())
@@ -43,6 +53,10 @@ mod test {
     use std::vec::Vec;
 
     async fn rocket() -> Rocket<Build> {
+        let max_connections = 5;
+        let min_connections = 1;
+        let idle_timeout = 15_000;
+
         let clusters = Clusters::new();
         clusters
             .add(-1, &pgml_dashboard::guards::default_database_url())
