fix db async session in process issue (#323)

Signed-off-by: Aisuko <[email protected]>

backend/src/api/dependencies/session.py

@@ -8,9 +8,13 @@
 async def get_async_session() -> typing.AsyncGenerator[SQLAlchemyAsyncSession, None]:
     try:
         yield async_db.async_session
+        loguru.logger.info(f"Async pool:{async_db.pool.status()}")
     except Exception as e:
         loguru.logger.info(f"Exception --- {e}")
         await async_db.async_session.rollback()
         raise
+    else:
+        # https://github.com/grillazz/fastapi-sqlalchemy-asyncpg/issues/63
+        await async_db.async_session.commit()
     finally:
         await async_db.async_session.close()

backend/src/api/routes/rag_datasets.py

@@ -20,6 +20,7 @@
 from src.models.schemas.dataset import RagDatasetCreate, RagDatasetResponse, LoadRAGDSResponse
 from src.repository.rag_datasets_eng import DatasetEng
 from src.repository.crud.account import AccountCRUDRepository
+from src.repository.crud.dataset_db import DataSetCRUDRepository
 from src.securities.authorizations.jwt import jwt_required
 from src.repository.crud.chat import SessionCRUDRepository
 from src.utilities.formatters.ds_formatter import DatasetFormatter
@@ -39,6 +40,7 @@
 )
 async def get_dataset_list(
     token: str = fastapi.Depends(oauth2_scheme),
+    ds_repo: DataSetCRUDRepository = fastapi.Depends(get_repository(repo_type=DataSetCRUDRepository)),
 ) -> list[RagDatasetResponse]:
     """
     Get all the pre-processed dataset list.
@@ -62,7 +64,7 @@ async def get_dataset_list(
 
     """
     # It is unthread safe
-    # list_ds = await ds_repo.get_dataset_list()
+    # list_ds_from_db = await ds_repo.get_dataset_list()
 
     list_ds = [settings.DEFAULT_RAG_DS_NAME]
 

backend/src/repository/database.py

@@ -20,7 +20,6 @@
     create_async_engine as create_sqlalchemy_async_engine,
 )
 from sqlalchemy.pool import Pool as SQLAlchemyPool
-from sqlalchemy.pool import NullPool
 from sqlalchemy import create_engine
 from src.config.manager import settings
 
@@ -33,13 +32,14 @@ def __init__(self):
         self.async_engine: SQLAlchemyAsyncEngine = create_sqlalchemy_async_engine(
             url=self.set_async_db_uri,
             echo=settings.IS_DB_ECHO_LOG,
+            echo_pool=True,
             future=True,
-            # pool_size=settings.DB_POOL_SIZE,
+            pool_size=settings.DB_POOL_SIZE,
+            pool_timeout=60,
             # max_overflow=settings.DB_POOL_OVERFLOW,
             # pool_recycle=3600,  # Periodically recycle connections (optional)
-            # pool_pre_ping=True,  # Check the connection status before using it
+            pool_pre_ping=True,  # Check the connection status before using it
             # https://github.com/MagicStack/asyncpg/issues/863
-            poolclass=NullPool,
         )
         self.sync_engine = create_engine(
             f"{settings.POSTGRES_SCHEMA}://{settings.POSTGRES_USERNAME}:{settings.POSTGRES_PASSWORD}@{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}"