🎨 优化数据库代码

2 files changed, 20 insertions, 73 deletions

diff --git a/ATRI/plugins/thesaurus/data_source.py b/ATRI/plugins/thesaurus/data_source.py
index 4f99b2c..281a777 100644
--- a/ATRI/plugins/thesaurus/data_source.py
+++ b/ATRI/plugins/thesaurus/data_source.py
@@ -2,23 +2,23 @@ from datetime import datetime, timedelta, timezone as tz
 
 from ATRI.message import MessageBuilder
 from ATRI.exceptions import ThesaurusError
+from ATRI.database import DatabaseWrapper, ThesaurusStoragor, ThesaurusAuditList
 
-from .db import DBForTS, DBForTAL
-from .db import ThesaurusStoragor
+
+DBForTS = DatabaseWrapper(ThesaurusStoragor)
+DBForTAL = DatabaseWrapper(ThesaurusAuditList)
 
 
 class ThesaurusManager:
     async def __add_item(self, _id: str, group_id: int, is_main: bool = False):
         if is_main:
             try:
-                async with DBForTS() as db:
-                    await db.add_item(_id, group_id)
+                await DBForTS.add_sub(_id=_id, group_id=group_id)
             except Exception:
                 raise ThesaurusError(f"添加词库(ts)数据失败 目标词id: {_id}")
         else:
             try:
-                async with DBForTAL() as db:
-                    await db.add_item(_id, group_id)
+                await DBForTAL.add_sub(_id=_id, group_id=group_id)
             except Exception:
                 raise ThesaurusError(f"添加词库(tal)数据失败 目标词id: {_id}")
 
@@ -27,56 +27,52 @@ class ThesaurusManager:
     ):
         if is_main:
             try:
-                async with DBForTS() as db:
-                    await db.update_item(_id, group_id, update_map)
+                await DBForTS.update_sub(
+                    update_map=update_map, _id=_id, group_id=group_id
+                )
             except Exception:
                 raise ThesaurusError(f"更新词库(ts)数据失败 目标词id: {_id}")
         else:
             try:
-                async with DBForTAL() as db:
-                    await db.update_item(_id, group_id, update_map)
+                await DBForTAL.update_sub(
+                    update_map=update_map, _id=_id, group_id=group_id
+                )
             except Exception:
                 raise ThesaurusError(f"更新词库(tal)数据失败 目标词id: {_id}")
 
     async def __del_item(self, _id: str, group_id: int, is_main: bool = False):
         if is_main:
             try:
-                async with DBForTS() as db:
-                    await db.del_item({"_id": _id, "group_id": group_id})
+                await DBForTS.del_sub({"_id": _id, "group_id": group_id})
             except Exception:
                 raise ThesaurusError(f"删除词库(ts)数据失败 目标词id: {_id}")
         else:
             try:
-                async with DBForTAL() as db:
-                    await db.del_item({"_id": _id, "group_id": group_id})
+                await DBForTAL.del_sub({"_id": _id, "group_id": group_id})
             except Exception:
                 raise ThesaurusError(f"删除词库(tal)数据失败 目标词id: {_id}")
 
     async def get_item_list(self, query_map: dict, is_main: bool = False) -> list:
         if is_main:
             try:
-                async with DBForTS() as db:
-                    return await db.get_item_list(query_map)
+                return await DBForTS.get_sub_list(query_map)
             except Exception:
                 raise ThesaurusError("获取词库(ts)列表数据失败")
         else:
             try:
-                async with DBForTAL() as db:
-                    return await db.get_item_list(query_map)
+                return await DBForTAL.get_sub_list(query_map)
             except Exception:
                 raise ThesaurusError("获取词库(tal)列表数据失败")
 
     async def get_all_items(self, is_main: bool = False) -> list:
         if is_main:
             try:
-                async with DBForTS() as db:
-                    return await db.get_all_items()
+                return await DBForTS.get_all_subs()
             except Exception:
                 raise ThesaurusError("获取全部词库(ts)列表数据失败")
         else:
             try:
-                async with DBForTAL() as db:
-                    return await db.get_all_items()
+                return await DBForTAL.get_all_subs()
             except Exception:
                 raise ThesaurusError("获取全部词库(tal)列表数据失败")
 
@@ -165,8 +161,7 @@ class ThesaurusManager:
 class ThesaurusListener:
     async def get_item_by_id(self, _id: str) -> ThesaurusStoragor:
         try:
-            async with DBForTS() as db:
-                data = await db.get_item_list({"_id": _id})
+            data = await DBForTS.get_sub_list({"_id": _id})
         except Exception:
             raise ThesaurusError(f"获取词库(ts)数据失败 词条ID: {_id}")
 
@@ -174,7 +169,6 @@ class ThesaurusListener:
 
     async def get_item_list(self, group_id: int):
         try:
-            async with DBForTS() as db:
-                return await db.get_item_list({"group_id": group_id})
+            return await DBForTS.get_sub_list({"group_id": group_id})
         except Exception:
             raise ThesaurusError(f"获取词库(ts)数据失败 目标群号: {group_id}")
diff --git a/ATRI/plugins/thesaurus/db.py b/ATRI/plugins/thesaurus/db.py
deleted file mode 100644
index b5394fc..0000000
--- a/ATRI/plugins/thesaurus/db.py
+++ /dev/null
@@ -1,47 +0,0 @@
-from ATRI.database import ThesaurusStoragor, ThesaurusAuditList
-
-
-class DBForTS:
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-    async def add_item(self, _id: str, group_id: int):
-        await ThesaurusStoragor.create(_id=_id, group_id=group_id)
-
-    async def update_item(self, _id: str, group_id: int, update_map: dict):
-        await ThesaurusStoragor.filter(_id=_id, group_id=group_id).update(**update_map)
-
-    async def del_item(self, query_map: dict):
-        await ThesaurusStoragor.filter(**query_map).delete()
-
-    async def get_item_list(self, query_map: dict) -> list:
-        return await ThesaurusStoragor.filter(**query_map)
-
-    async def get_all_items(self) -> list:
-        return await ThesaurusStoragor.all()
-
-
-class DBForTAL:
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-    async def add_item(self, _id: str, group_id: int):
-        await ThesaurusAuditList.create(_id=_id, group_id=group_id)
-
-    async def update_item(self, _id: str, group_id: int, update_map: dict):
-        await ThesaurusAuditList.filter(_id=_id, group_id=group_id).update(**update_map)
-
-    async def del_item(self, query_map: dict):
-        await ThesaurusAuditList.filter(**query_map).delete()
-
-    async def get_item_list(self, query_map: dict) -> list:
-        return await ThesaurusAuditList.filter(**query_map)
-
-    async def get_all_items(self) -> list:
-        return await ThesaurusAuditList.all()