From 73dfd2a849c52e14a65eeda7b18d0ce29c120562 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9C=B1=E5=A4=A7=E6=B4=8B?= <714403855@qq.com>
Date: Fri, 12 Dec 2025 15:04:55 +0800
Subject: [PATCH] add polardb pool config

---
 src/memos/api/config.py        | 1 +
 src/memos/configs/graph_db.py  | 4 ++++
 src/memos/graph_dbs/polardb.py | 6 +++++-
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/memos/api/config.py b/src/memos/api/config.py
index af0f0473d..9aa4dba5d 100644
--- a/src/memos/api/config.py
+++ b/src/memos/api/config.py
@@ -590,6 +590,7 @@ def get_polardb_config(user_id: str | None = None) -> dict[str, Any]:
             "user": os.getenv("POLAR_DB_USER", "root"),
             "password": os.getenv("POLAR_DB_PASSWORD", "123456"),
             "db_name": db_name,
+            "maxconn": int(os.getenv("POLARDB_POOL_MAX_CONN", "100")),
             "user_name": user_name,
             "use_multi_db": use_multi_db,
             "auto_create": True,
diff --git a/src/memos/configs/graph_db.py b/src/memos/configs/graph_db.py
index ce180606b..3b4bace0e 100644
--- a/src/memos/configs/graph_db.py
+++ b/src/memos/configs/graph_db.py
@@ -198,6 +198,10 @@ class PolarDBGraphDBConfig(BaseConfig):
         ),
     )
     embedding_dimension: int = Field(default=1024, description="Dimension of vector embedding")
+    maxconn: int = Field(
+        default=100,
+        description="Maximum number of connections in the connection pool",
+    )
 
     @model_validator(mode="after")
     def validate_config(self):
diff --git a/src/memos/graph_dbs/polardb.py b/src/memos/graph_dbs/polardb.py
index 50ff4ab90..c3f0297b7 100644
--- a/src/memos/graph_dbs/polardb.py
+++ b/src/memos/graph_dbs/polardb.py
@@ -136,6 +136,7 @@ def __init__(self, config: PolarDBGraphDBConfig):
             port = config.get("port")
             user = config.get("user")
             password = config.get("password")
+            maxconn = config.get("maxconn", 100)  # De
         else:
             self.db_name = config.db_name
             self.user_name = config.user_name
@@ -143,17 +144,19 @@ def __init__(self, config: PolarDBGraphDBConfig):
             port = config.port
             user = config.user
             password = config.password
+            maxconn = config.maxconn if hasattr(config, "maxconn") else 100
         """
         # Create connection
         self.connection = psycopg2.connect(
             host=host, port=port, user=user, password=password, dbname=self.db_name,minconn=10, maxconn=2000
         )
         """
+        logger.info(f" db_name: {self.db_name} current maxconn is:'{maxconn}'")
 
         # Create connection pool
         self.connection_pool = psycopg2.pool.ThreadedConnectionPool(
             minconn=5,
-            maxconn=100,
+            maxconn=maxconn,
             host=host,
             port=port,
             user=user,
@@ -216,6 +219,7 @@ def _get_connection(self):
         Raises:
             RuntimeError: If connection pool is closed or exhausted after retries
         """
+        logger.info(f" db_name: {self.db_name} pool maxconn is:'{self.connection_pool.maxconn}'")
         if self._pool_closed:
             raise RuntimeError("Connection pool has been closed")