diff --git a/src/database/database.py b/src/database/database.py index 12dbd40..fbdc4b6 100644 --- a/src/database/database.py +++ b/src/database/database.py @@ -67,6 +67,14 @@ class DatabaseManager: self.is_connected = False logger.info("数据库连接已断开") + async def session(self): + """获取数据库会话的异步上下文管理器""" + if not self.is_connected: + await self.connect() + + async with self.async_session() as session: + yield session + async def get_session(self) -> AsyncSession: """获取数据库会话""" if not self.is_connected: diff --git a/src/database/init_db.py b/src/database/init_db.py index 8e9bed8..120aa67 100644 --- a/src/database/init_db.py +++ b/src/database/init_db.py @@ -119,8 +119,13 @@ async def init_database(): await db_manager.connect() await db_manager.create_tables() - async with db_manager.session() as session: + async with db_manager.get_session() as session: perm_map = await init_permissions(session) + if perm_map is None: + # Permissions already existed, fetch them from database + result = await session.execute(select(Permission)) + perms = result.scalars().all() + perm_map = {perm.code: perm.id for perm in perms} await init_roles(session, perm_map) await create_admin_user(session)