feat(conn-pool): tweak connection pool options and expose min/max conn (#851)

* feat(conn-pool): tweak connection pool options and expose min/max conn

* docs(conn-pool): new options about min/max connections

Author: georgeh0

docs/docs/core/settings.mdx

@@ -77,24 +77,33 @@ If not set, all flows are in a default unnamed namespace.
 
 `DatabaseConnectionSpec` configures the connection to a database. Only Postgres is supported for now. It has the following fields:
 
-*   `url` (type: `str`, required): The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`.
+*   `url` (type: `str`): The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`.
 
     *Environment variable* for `Settings.database.url`: `COCOINDEX_DATABASE_URL`
 
-*   `user` (type: `str`, optional): The username for the Postgres database. If not provided, username will come from `url`.
+*   `user` (type: `str | None`, default: `None`): The username for the Postgres database. If not provided, username will come from `url`.
 
     *Environment variable* for `Settings.database.user`: `COCOINDEX_DATABASE_USER`
 
-*   `password` (type: `str`, optional): The password for the Postgres database. If not provided, password will come from `url`.
+*   `password` (type: `str | None`, default: `None`): The password for the Postgres database. If not provided, password will come from `url`.
 
     *Environment variable* for `Settings.database.password`: `COCOINDEX_DATABASE_PASSWORD`
 
-:::tip
+    :::tip
 
-Please be careful that all values in `url` needs to be url-encoded if they contain special characters.
-For this reason, prefer to use the separated `user` and `password` fields for username and password.
+    Please be careful that all values in `url` needs to be url-encoded if they contain special characters.
+    For this reason, prefer to use the separated `user` and `password` fields for username and password.
+
+    :::
+
+*   `max_connections` (type: `int`, default: `64`): The maximum number of connections to keep in the pool.
+
+    *Environment variable* for `Settings.database.max_connections`: `COCOINDEX_DATABASE_MAX_CONNECTIONS`
+
+*   `min_connections` (type: `int`, default: `16`): The minimum number of connections to keep in the pool.
+
+    *Environment variable* for `Settings.database.min_connections`: `COCOINDEX_DATABASE_MIN_CONNECTIONS`
 
-:::
 
 :::info
 
@@ -125,5 +134,7 @@ This is the list of environment variables, each of which has a corresponding fie
 | `COCOINDEX_DATABASE_URL` | `database.url` | Yes |
 | `COCOINDEX_DATABASE_USER` | `database.user` | No |
 | `COCOINDEX_DATABASE_PASSWORD` | `database.password` | No |
+| `COCOINDEX_DATABASE_MAX_CONNECTIONS` | `database.max_connections` | No (default: `64`) |
+| `COCOINDEX_DATABASE_MIN_CONNECTIONS` | `database.min_connections` | No (default: `16`) |
 | `COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS` | `global_execution_options.source_max_inflight_rows` | No (default: `256`) |
 | `COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES` | `global_execution_options.source_max_inflight_bytes` | No |

python/cocoindex/setting.py

@@ -44,6 +44,8 @@ class DatabaseConnectionSpec:
     url: str
     user: str | None = None
     password: str | None = None
+    max_connections: int = 64
+    min_connections: int = 16
 
 
 @dataclass
@@ -92,10 +94,22 @@ def from_env(cls) -> Self:
 
         database_url = os.getenv("COCOINDEX_DATABASE_URL")
         if database_url is not None:
-            db_kwargs: dict[str, str] = dict()
+            db_kwargs: dict[str, Any] = dict()
             _load_field(db_kwargs, "url", "COCOINDEX_DATABASE_URL", required=True)
             _load_field(db_kwargs, "user", "COCOINDEX_DATABASE_USER")
             _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
+            _load_field(
+                db_kwargs,
+                "max_connections",
+                "COCOINDEX_DATABASE_MAX_CONNECTIONS",
+                parse=int,
+            )
+            _load_field(
+                db_kwargs,
+                "min_connections",
+                "COCOINDEX_DATABASE_MIN_CONNECTIONS",
+                parse=int,
+            )
             database = DatabaseConnectionSpec(**db_kwargs)
         else:
             database = None

src/lib_context.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use crate::prelude::*;
 
 use crate::builder::AnalyzedFlow;
@@ -7,7 +9,7 @@ use crate::settings;
 use crate::setup::ObjectSetupStatus;
 use axum::http::StatusCode;
 use sqlx::PgPool;
-use sqlx::postgres::PgConnectOptions;
+use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
 use tokio::runtime::Runtime;
 
 pub struct FlowExecutionContext {
@@ -176,7 +178,29 @@ impl DbPools {
                 if let Some(password) = &conn_spec.password {
                     pg_options = pg_options.password(password);
                 }
-                let pool = PgPool::connect_with(pg_options)
+
+                // Try to connect to the database with a low timeout first.
+                {
+                    let pool_options = PgPoolOptions::new()
+                        .max_connections(1)
+                        .min_connections(1)
+                        .acquire_timeout(Duration::from_secs(30));
+                    let pool = pool_options
+                        .connect_with(pg_options.clone())
+                        .await
+                        .context(format!("Failed to connect to database {}", conn_spec.url))?;
+                    let _ = pool.acquire().await?;
+                }
+
+                // Now create the actual pool.
+                let pool_options = PgPoolOptions::new()
+                    .max_connections(conn_spec.max_connections)
+                    .min_connections(conn_spec.min_connections)
+                    .acquire_timeout(Duration::from_secs(5 * 60))
+                    .idle_timeout(Duration::from_secs(10 * 60))
+                    .max_lifetime(Duration::from_secs(60 * 60));
+                let pool = pool_options
+                    .connect_with(pg_options)
                     .await
                     .context("Failed to connect to database")?;
                 anyhow::Ok(pool)
@@ -330,6 +354,8 @@ mod tests {
                 url: "postgresql://test".to_string(),
                 user: None,
                 password: None,
+                max_connections: 10,
+                min_connections: 1,
             }),
             ..Default::default()
         };

src/settings.rs

@@ -5,6 +5,8 @@ pub struct DatabaseConnectionSpec {
     pub url: String,
     pub user: Option<String>,
     pub password: Option<String>,
+    pub max_connections: u32,
+    pub min_connections: u32,
 }
 
 #[derive(Deserialize, Debug, Default)]
@@ -34,7 +36,9 @@ mod tests {
             "database": {
                 "url": "postgresql://localhost:5432/test",
                 "user": "testuser",
-                "password": "testpass"
+                "password": "testpass",
+                "min_connections": 1,
+                "max_connections": 10
             },
             "app_namespace": "test_app"
         }"#;
@@ -46,6 +50,8 @@ mod tests {
         assert_eq!(db.url, "postgresql://localhost:5432/test");
         assert_eq!(db.user, Some("testuser".to_string()));
         assert_eq!(db.password, Some("testpass".to_string()));
+        assert_eq!(db.min_connections, 1);
+        assert_eq!(db.max_connections, 10);
         assert_eq!(settings.app_namespace, "test_app");
     }
 
@@ -75,7 +81,9 @@ mod tests {
     fn test_settings_deserialize_database_without_user_password() {
         let json = r#"{
             "database": {
-                "url": "postgresql://localhost:5432/test"
+                "url": "postgresql://localhost:5432/test",
+                "min_connections": 1,
+                "max_connections": 10
             }
         }"#;
 
@@ -86,6 +94,8 @@ mod tests {
         assert_eq!(db.url, "postgresql://localhost:5432/test");
         assert_eq!(db.user, None);
         assert_eq!(db.password, None);
+        assert_eq!(db.min_connections, 1);
+        assert_eq!(db.max_connections, 10);
         assert_eq!(settings.app_namespace, "");
     }
 
@@ -94,13 +104,17 @@ mod tests {
         let json = r#"{
             "url": "postgresql://localhost:5432/test",
             "user": "testuser",
-            "password": "testpass"
+            "password": "testpass",
+            "min_connections": 1,
+            "max_connections": 10
         }"#;
 
         let db_spec: DatabaseConnectionSpec = serde_json::from_str(json).unwrap();
 
         assert_eq!(db_spec.url, "postgresql://localhost:5432/test");
         assert_eq!(db_spec.user, Some("testuser".to_string()));
         assert_eq!(db_spec.password, Some("testpass".to_string()));
+        assert_eq!(db_spec.min_connections, 1);
+        assert_eq!(db_spec.max_connections, 10);
     }
 }