feat: add datasource timeout

backend/apps/datasource/models/datasource.py

@@ -73,6 +73,7 @@ class DatasourceConf(BaseModel):
     filename: str = ''
     sheets: List = ''
     mode: str = ''
+    timeout: int = 30
 
     def to_dict(self):
         return {
@@ -86,7 +87,8 @@ class DatasourceConf(BaseModel):
             "dbSchema": self.dbSchema,
             "filename": self.filename,
             "sheets": self.sheets,
-            "mode": self.mode
+            "mode": self.mode,
+            "timeout": self.timeout
         }
 
 

backend/apps/db/db.py

@@ -48,11 +48,17 @@ def get_uri(ds: CoreDatasource):
 
 def get_engine(ds: CoreDatasource) -> Engine:
     conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
+    if conf.timeout is None:
+        conf.timeout = 30
     if ds.type == "pg" and (conf.dbSchema is not None and conf.dbSchema != ""):
-        engine = create_engine(get_uri(ds), connect_args={"options": f"-c search_path={conf.dbSchema}"},
-                               pool_timeout=30, pool_size=20, max_overflow=10)
+        engine = create_engine(get_uri(ds),
+                               connect_args={"options": f"-c search_path={conf.dbSchema}",
+                                             "connect_timeout": conf.timeout},
+                               pool_timeout=conf.timeout, pool_size=20, max_overflow=10)
     else:
-        engine = create_engine(get_uri(ds), pool_timeout=30, pool_size=20, max_overflow=10)
+        engine = create_engine(get_uri(ds), connect_args={"connect_timeout": conf.timeout}, pool_timeout=conf.timeout,
+                               pool_size=20,
+                               max_overflow=10)
     return engine
 
 

backend/apps/db/engine.py

@@ -13,7 +13,7 @@ from common.core.config import settings
 def get_engine_config():
     return DatasourceConf(username=settings.POSTGRES_USER, password=settings.POSTGRES_PASSWORD,
                           host=settings.POSTGRES_SERVER, port=settings.POSTGRES_PORT, database=settings.POSTGRES_DB,
-                          dbSchema="public")
+                          dbSchema="public", timeout=30)
 
 
 def get_engine_uri(conf: DatasourceConf):