improvement: add ARGILLA_DATABASE_POSTGRESQL_POOL_SIZE and ARGILLA_DATABASE_POSTGRESQL_MAX_OVERFLOW (#5220)

# Description

After testing a high number of concurrent requests using PostgreSQL I
received the following error:

```
QueuePool limit of size 5 overflow 10 reached, connection timed out, timeout 30.00
```

This PR add the following two environment variables so we can configure
the pool size and max overflow.

Refs #5000 

**Type of change**

- Improvement (change adding some improvement to an existing
functionality)

**How Has This Been Tested**

- [x] Manually testing with PostgreSQL.

**Checklist**

- I added relevant documentation
- I followed the style guidelines of this project
- I did a self-review of my code
- I made corresponding changes to the documentation
- I confirm My changes generate no new warnings
- I have added tests that prove my fix is effective or that my feature
works
- I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

Author: jfcalvo

argilla-server/CHANGELOG.md

@@ -19,6 +19,8 @@ These are the section headers that we use:
 ### Added
 
 - Added new `ARGILLA_DATABASE_SQLITE_TIMEOUT` environment variable allowing to set transactions timeout for SQLite. ([#5213](https://github.com/argilla-io/argilla/pull/5213))
+- Added new `ARGILLA_DATABASE_POSTGRESQL_POOL_SIZE` environment variable allowing to set the number of connections to keep open inside the database connection pool. ([#5220](https://github.com/argilla-io/argilla/pull/5220))
+- Added new `ARGILLA_DATABASE_POSTGRESQL_MAX_OVERFLOW` environment variable allowing to set the number of connections that can be opened above and beyond the `ARGILLA_DATABASE_POSTGRESQL_POOL_SIZE` setting. ([#5220](https://github.com/argilla-io/argilla/pull/5220))
 
 ### Fixed
 

argilla-server/src/argilla_server/constants.py

@@ -16,6 +16,7 @@
 WORKSPACE_HEADER_NAME = "X-Argilla-Workspace"
 
 DATABASE_SQLITE = "sqlite"
+DATABASE_POSTGRESQL = "postgresql"
 
 SEARCH_ENGINE_ELASTICSEARCH = "elasticsearch"
 SEARCH_ENGINE_OPENSEARCH = "opensearch"
@@ -26,6 +27,9 @@
 
 DEFAULT_DATABASE_SQLITE_TIMEOUT = 15
 
+DEFAULT_DATABASE_POSTGRESQL_POOL_SIZE = 15
+DEFAULT_DATABASE_POSTGRESQL_MAX_OVERFLOW = 10
+
 DEFAULT_MAX_KEYWORD_LENGTH = 128
 DEFAULT_TELEMETRY_KEY = "C6FkcaoCbt78rACAgvyBxGBcMB3dM3nn"
 

argilla-server/src/argilla_server/database.py

@@ -51,7 +51,8 @@ def set_sqlite_pragma(dbapi_connection, connection_record):
         cursor.close()
 
 
-async_engine = create_async_engine(settings.database_url, connect_args=settings.database_connect_args)
+async_engine = create_async_engine(settings.database_url, **settings.database_engine_args)
+
 AsyncSessionLocal = async_sessionmaker(autocommit=False, expire_on_commit=False, bind=async_engine)
 
 

argilla-server/src/argilla_server/settings.py

@@ -27,13 +27,16 @@
 
 from argilla_server.constants import (
     DATABASE_SQLITE,
+    DATABASE_POSTGRESQL,
     DEFAULT_LABEL_SELECTION_OPTIONS_MAX_ITEMS,
     DEFAULT_MAX_KEYWORD_LENGTH,
     DEFAULT_SPAN_OPTIONS_MAX_ITEMS,
     DEFAULT_TELEMETRY_KEY,
     DEFAULT_DATABASE_SQLITE_TIMEOUT,
     SEARCH_ENGINE_ELASTICSEARCH,
     SEARCH_ENGINE_OPENSEARCH,
+    DEFAULT_DATABASE_POSTGRESQL_POOL_SIZE,
+    DEFAULT_DATABASE_POSTGRESQL_MAX_OVERFLOW,
 )
 from argilla_server.pydantic_v1 import BaseSettings, Field, root_validator, validator
 
@@ -75,7 +78,19 @@ class Settings(BaseSettings):
 
     home_path: Optional[str] = Field(description="The home path where argilla related files will be stored")
     base_url: Optional[str] = Field(description="The default base url where server will be deployed")
+
     database_url: Optional[str] = Field(description="The database url that argilla will use as data store")
+    # https://docs.sqlalchemy.org/en/20/core/engines.html#sqlalchemy.create_engine.params.pool_size
+    database_postgresql_pool_size: Optional[int] = Field(
+        default=DEFAULT_DATABASE_POSTGRESQL_POOL_SIZE,
+        description="The number of connections to keep open inside the database connection pool",
+    )
+    # https://docs.sqlalchemy.org/en/20/core/engines.html#sqlalchemy.create_engine.params.max_overflow
+    database_postgresql_max_overflow: Optional[int] = Field(
+        default=DEFAULT_DATABASE_POSTGRESQL_MAX_OVERFLOW,
+        description="The number of connections that can be opened above and beyond the pool_size setting",
+    )
+    # https://docs.python.org/3/library/sqlite3.html#sqlite3.connect
     database_sqlite_timeout: Optional[int] = Field(
         default=DEFAULT_DATABASE_SQLITE_TIMEOUT,
         description="SQLite database connection timeout in seconds",
@@ -227,9 +242,19 @@ def old_dataset_records_index_name(self) -> str:
         return index_name.replace("<NAMESPACE>", f".{ns}")
 
     @property
-    def database_connect_args(self) -> Dict:
+    def database_engine_args(self) -> Dict:
         if self.database_is_sqlite:
-            return {"timeout": self.database_sqlite_timeout}
+            return {
+                "connect_args": {
+                    "timeout": self.database_sqlite_timeout,
+                },
+            }
+
+        if self.database_is_postgresql:
+            return {
+                "pool_size": self.database_postgresql_pool_size,
+                "max_overflow": self.database_postgresql_max_overflow,
+            }
 
         return {}
 
@@ -240,6 +265,13 @@ def database_is_sqlite(self) -> bool:
 
         return self.database_url.lower().startswith(DATABASE_SQLITE)
 
+    @property
+    def database_is_postgresql(self) -> bool:
+        if self.database_url is None:
+            return False
+
+        return self.database_url.lower().startswith(DATABASE_POSTGRESQL)
+
     @property
     def search_engine_is_elasticsearch(self) -> bool:
         return self.search_engine == SEARCH_ENGINE_ELASTICSEARCH

argilla-server/tests/unit/commons/test_settings.py

@@ -59,16 +59,6 @@ def test_settings_default_database_url(monkeypatch):
     assert settings.database_url == f"sqlite+aiosqlite:///{settings.home_path}/argilla.db?check_same_thread=False"
 
 
-def test_settings_database_sqlite_timeout(monkeypatch):
-    monkeypatch.setenv("ARGILLA_DATABASE_SQLITE_TIMEOUT", "3")
-
-    assert Settings().database_sqlite_timeout == 3
-
-
-def test_settings_default_database_sqlite_timeout():
-    assert Settings().database_sqlite_timeout == 15
-
-
 @pytest.mark.parametrize(
     "url, expected_url",
     [
@@ -82,3 +72,33 @@ def test_settings_database_url(url: str, expected_url: str, monkeypatch):
     monkeypatch.setenv("ARGILLA_DATABASE_URL", url)
 
     assert Settings().database_url == expected_url
+
+
+def test_settings_default_database_sqlite_timeout():
+    assert Settings().database_sqlite_timeout == 15
+
+
+def test_settings_database_sqlite_timeout(monkeypatch):
+    monkeypatch.setenv("ARGILLA_DATABASE_SQLITE_TIMEOUT", "3")
+
+    assert Settings().database_sqlite_timeout == 3
+
+
+def test_settings_default_database_postgresql_pool_size():
+    assert Settings().database_postgresql_pool_size == 15
+
+
+def test_settings_database_postgresql_pool_size(monkeypatch):
+    monkeypatch.setenv("ARGILLA_DATABASE_POSTGRESQL_POOL_SIZE", "42")
+
+    assert Settings().database_postgresql_pool_size == 42
+
+
+def test_settings_default_database_postgresql_max_overflow():
+    assert Settings().database_postgresql_max_overflow == 10
+
+
+def test_settings_database_postgresql_max_overflow(monkeypatch):
+    monkeypatch.setenv("ARGILLA_DATABASE_POSTGRESQL_MAX_OVERFLOW", "12")
+
+    assert Settings().database_postgresql_max_overflow == 12

docs/_source/getting_started/installation/configurations/server_configuration.md

@@ -73,6 +73,14 @@ The following environment variables are useful only when SQLite is used:
 
 - `ARGILLA_DATABASE_SQLITE_TIMEOUT`: How many seconds the connection should wait before raising an `OperationalError` when a table is locked. If another connection opens a transaction to modify a table, that table will be locked until the transaction is committed. (Defaut: `15` seconds).
 
+#### PostgreSQL
+
+The following environment variables are useful only when PostgreSQL is used:
+
+- `ARGILLA_DATABASE_POSTGRESQL_POOL_SIZE`: The number of connections to keep open inside the database connection pool (Default: `15`).
+
+- `ARGILLA_DATABASE_POSTGRESQL_MAX_OVERFLOW`: The number of connections that can be opened above and beyond `ARGILLA_DATABASE_POSTGRESQL_POOL_SIZE` setting (Default: `10`).
+
 #### Elasticsearch and Opensearch
 
 - `ARGILLA_ELASTICSEARCH`: URL of the connection endpoint of the Elasticsearch instance (Default: `http://localhost:9200`).