From 314b04eabaef6dc59d63dbdcf727b8b9ebfb0a63 Mon Sep 17 00:00:00 2001
From: Zamil Majdy <zamil.majdy@agpt.co>
Date: Tue, 31 Dec 2024 23:19:15 +0700
Subject: [PATCH] feat(backend): Make scheduler DB connection pool configurable
 & prevent connection overflow (#9149)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

max_overflow parameter description:
```
    :param max_overflow=10: the number of connections to allow in
        connection pool "overflow", that is connections that can be
        opened above and beyond the pool_size setting, which defaults
        to five. this is only used with :class:`~sqlalchemy.pool.QueuePool`.
```

### Changes 🏗️

* Prevent additional db connections from being created in addition to
the pool size for the scheduler.
* Make the pool size configurable.

### Checklist 📋

#### For code changes:
- [ ] I have clearly listed my changes in the PR description
- [ ] I have made a test plan
- [ ] I have tested my changes according to the test plan:
  <!-- Put your test plan here: -->
  - [ ] ...

<details>
  <summary>Example test plan</summary>

  - [ ] Create from scratch and execute an agent with at least 3 blocks
- [ ] Import an agent from file upload, and confirm it executes
correctly
  - [ ] Upload agent to marketplace
- [ ] Import an agent from marketplace and confirm it executes correctly
  - [ ] Edit an agent from monitor, and confirm it executes correctly
</details>

#### For configuration changes:
- [ ] `.env.example` is updated or already compatible with my changes
- [ ] `docker-compose.yml` is updated or already compatible with my
changes
- [ ] I have included a list of my configuration changes in the PR
description (under **Changes**)

<details>
  <summary>Examples of configuration changes</summary>

  - Changing ports
  - Adding new services that need to communicate with each other
  - Secrets or environment variable changes
  - New or infrastructure changes such as databases
</details>

Co-authored-by: Swifty <craigswift13@gmail.com>
---
 autogpt_platform/backend/backend/executor/scheduler.py | 10 +++++++++-
 autogpt_platform/backend/backend/util/settings.py      |  5 +++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/autogpt_platform/backend/backend/executor/scheduler.py b/autogpt_platform/backend/backend/executor/scheduler.py
index 3c906a3af..edf126f69 100644
--- a/autogpt_platform/backend/backend/executor/scheduler.py
+++ b/autogpt_platform/backend/backend/executor/scheduler.py
@@ -99,6 +99,10 @@ class ExecutionScheduler(AppService):
     def get_port(cls) -> int:
         return config.execution_scheduler_port
 
+    @classmethod
+    def db_pool_size(cls) -> int:
+        return config.scheduler_db_pool_size
+
     @property
     @thread_cached
     def execution_client(self) -> ExecutionManager:
@@ -110,7 +114,11 @@ class ExecutionScheduler(AppService):
         self.scheduler = BlockingScheduler(
             jobstores={
                 "default": SQLAlchemyJobStore(
-                    engine=create_engine(db_url),
+                    engine=create_engine(
+                        url=db_url,
+                        pool_size=self.db_pool_size(),
+                        max_overflow=0,
+                    ),
                     metadata=MetaData(schema=db_schema),
                 )
             }
diff --git a/autogpt_platform/backend/backend/util/settings.py b/autogpt_platform/backend/backend/util/settings.py
index 69504f528..409a18d93 100644
--- a/autogpt_platform/backend/backend/util/settings.py
+++ b/autogpt_platform/backend/backend/util/settings.py
@@ -153,6 +153,11 @@ class Config(UpdateTrackingModel["Config"], BaseSettings):
         description="The name of the Google Cloud Storage bucket for media files",
     )
 
+    scheduler_db_pool_size: int = Field(
+        default=3,
+        description="The pool size for the scheduler database connection pool",
+    )
+
     @field_validator("platform_base_url", "frontend_base_url")
     @classmethod
     def validate_platform_base_url(cls, v: str, info: ValidationInfo) -> str: