🎨 优化数据库代码

2 files changed, 8 insertions, 37 deletions

diff --git a/ATRI/plugins/rss/rss_rsshub/data_source.py b/ATRI/plugins/rss/rss_rsshub/data_source.py
index 9ca1c04..d829f1a 100644
--- a/ATRI/plugins/rss/rss_rsshub/data_source.py
+++ b/ATRI/plugins/rss/rss_rsshub/data_source.py
@@ -2,43 +2,40 @@ import xmltodict
 
 from ATRI.exceptions import RssError
 from ATRI.utils import request, gen_random_str
+from ATRI.database import DatabaseWrapper, RssRsshubSubcription
 
-from .db import DB
+
+DB = DatabaseWrapper(RssRsshubSubcription)
 
 
 class RssHubSubscriptor:
     async def __add_sub(self, _id: str, group_id: int):
         try:
-            async with DB() as db:
-                await db.add_sub(_id, group_id)
+            await DB.add_sub(_id=_id, group_id=group_id)
         except Exception:
             raise RssError("rss.rsshub: 添加订阅失败")
 
     async def update_sub(self, _id: str, group_id: int, update_map: dict):
         try:
-            async with DB() as db:
-                await db.update_sub(_id, group_id, update_map)
+            await DB.update_sub(update_map=update_map, _id=_id, group_id=group_id)
         except Exception:
             raise RssError("rss.rsshub: 更新订阅失败")
 
     async def __del_sub(self, _id: str, group_id: int):
         try:
-            async with DB() as db:
-                await db.del_sub({"_id": _id, "group_id": group_id})
+            await DB.del_sub({"_id": _id, "group_id": group_id})
         except Exception:
             raise RssError("rss.rsshub: 删除订阅失败")
 
     async def get_sub_list(self, query_map: dict) -> list:
         try:
-            async with DB() as db:
-                return await db.get_sub_list(query_map)
+            return await DB.get_sub_list(query_map)
         except Exception:
             raise RssError("rss.rsshub: 获取订阅列表失败")
 
     async def get_all_subs(self) -> list:
         try:
-            async with DB() as db:
-                return await db.get_all_subs()
+            return await DB.get_all_subs()
         except Exception:
             raise RssError("rss.rsshub: 获取所有订阅失败")
 
diff --git a/ATRI/plugins/rss/rss_rsshub/db.py b/ATRI/plugins/rss/rss_rsshub/db.py
deleted file mode 100644
index 3f614a3..0000000
--- a/ATRI/plugins/rss/rss_rsshub/db.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from ATRI.database import RssRsshubSubcription
-
-
-class DB:
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-    async def add_sub(self, _id: str, group_id: int):
-        await RssRsshubSubcription.create(_id=_id, group_id=group_id)
-
-    async def update_sub(self, _id, group_id, update_map: dict):
-        await RssRsshubSubcription.filter(_id=_id, group_id=group_id).update(
-            **update_map
-        )
-
-    async def del_sub(self, query_map: dict):
-        await RssRsshubSubcription.filter(**query_map).delete()
-
-    async def get_sub_list(self, query_map: dict) -> list:
-        return await RssRsshubSubcription.filter(**query_map)
-
-    async def get_all_subs(self) -> list:
-        return await RssRsshubSubcription.all()