🚚💩 移动数据库相关函数, 优化代码

2 files changed, 10 insertions, 4 deletions

diff --git a/ATRI/database/__init__.py b/ATRI/database/__init__.py
index 6c58d20..497bea4 100644
--- a/ATRI/database/__init__.py
+++ b/ATRI/database/__init__.py
@@ -1,2 +1,2 @@
-from .db import init_database
+from .db import init_database, close_database_connection
 from .models import BilibiliSubscription, TwitterSubscription
diff --git a/ATRI/database/db.py b/ATRI/database/db.py
index 6e20ad3..c8782de 100644
--- a/ATRI/database/db.py
+++ b/ATRI/database/db.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from tortoise import Tortoise, run_async
+from tortoise import Tortoise
 
 from ATRI.log import logger as log
 
@@ -21,7 +21,13 @@ async def run():
     await Tortoise.generate_schemas()
 
 
-def init_database():
+async def init_database():
     log.info("正在初始化数据库...")
-    run_async(run())
+    await run()
     log.success("数据库初始化完成")
+
+
+async def close_database_connection():
+    log.info("正在关闭数据库连接...")
+    await Tortoise.close_connections()
+    log.info("数据库成功关闭")