summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--ATRI/database/__init__.py1
-rw-r--r--ATRI/database/wrapper.py31
-rw-r--r--ATRI/plugins/bilibili_dynamic/data_source.py18
-rw-r--r--ATRI/plugins/bilibili_dynamic/db.py26
-rw-r--r--ATRI/plugins/rss/rss_mikanan/data_source.py18
-rw-r--r--ATRI/plugins/rss/rss_mikanan/db.py26
-rw-r--r--ATRI/plugins/rss/rss_rsshub/data_source.py19
-rw-r--r--ATRI/plugins/rss/rss_rsshub/db.py26
-rw-r--r--ATRI/plugins/thesaurus/data_source.py46
-rw-r--r--ATRI/plugins/thesaurus/db.py47
10 files changed, 74 insertions, 184 deletions
diff --git a/ATRI/database/__init__.py b/ATRI/database/__init__.py
index 5c12c6a..f43b83b 100644
--- a/ATRI/database/__init__.py
+++ b/ATRI/database/__init__.py
@@ -1,2 +1,3 @@
from .db import init_database, close_database_connection
from .models import *
+from .wrapper import DatabaseWrapper
diff --git a/ATRI/database/wrapper.py b/ATRI/database/wrapper.py
new file mode 100644
index 0000000..7838bd6
--- /dev/null
+++ b/ATRI/database/wrapper.py
@@ -0,0 +1,31 @@
+from typing import Type
+
+from tortoise.models import Model
+
+
+class DatabaseWrapper:
+ def __init__(self, model: Type[Model]):
+ self.model = model
+
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
+ pass
+
+ async def add_sub(self, *args, **kwargs):
+ await self.model.create(*args, **kwargs)
+
+ async def update_sub(self, update_map: dict, **kwargs):
+ await self.model.filter(**kwargs).update(
+ **update_map
+ )
+
+ async def del_sub(self, query_map: dict):
+ await self.model.filter(**query_map).delete()
+
+ async def get_sub_list(self, query_map: dict) -> list:
+ return await self.model.filter(**query_map)
+
+ async def get_all_subs(self) -> list:
+ return await self.model.all()
diff --git a/ATRI/plugins/bilibili_dynamic/data_source.py b/ATRI/plugins/bilibili_dynamic/data_source.py
index d629454..15babe5 100644
--- a/ATRI/plugins/bilibili_dynamic/data_source.py
+++ b/ATRI/plugins/bilibili_dynamic/data_source.py
@@ -5,8 +5,8 @@ from operator import itemgetter
from ATRI.message import MessageBuilder
from ATRI.utils import TimeDealer
from ATRI.exceptions import BilibiliDynamicError
+from ATRI.database import DatabaseWrapper, BilibiliSubscription
-from .db import DB
from .api import API
@@ -17,27 +17,25 @@ _OUTPUT_FORMAT = (
.text("链接: {up_dy_link}")
.done()
)
+DB = DatabaseWrapper(BilibiliSubscription)
class BilibiliDynamicSubscriptor:
async def __add_sub(self, uid: int, group_id: int):
try:
- async with DB() as db:
- await db.add_sub(uid, group_id)
+ await DB.add_sub(uid=uid, group_id=group_id)
except Exception:
raise BilibiliDynamicError("添加订阅失败")
async def update_sub(self, uid: int, group_id: int, update_map: dict):
try:
- async with DB() as db:
- await db.update_sub(uid, group_id, update_map)
+ await DB.update_sub(update_map=update_map, uid=uid, group_id=group_id)
except Exception:
raise BilibiliDynamicError("更新订阅失败")
async def __del_sub(self, uid: int, group_id: int):
try:
- async with DB() as db:
- await db.del_sub({"uid": uid, "group_id": group_id})
+ await DB.del_sub({"uid": uid, "group_id": group_id})
except Exception:
raise BilibiliDynamicError("删除订阅失败")
@@ -48,15 +46,13 @@ class BilibiliDynamicSubscriptor:
query_map = {"uid": uid, "group_id": group_id}
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 BilibiliDynamicError("获取订阅列表失败")
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 BilibiliDynamicError("获取全部订阅列表失败")
diff --git a/ATRI/plugins/bilibili_dynamic/db.py b/ATRI/plugins/bilibili_dynamic/db.py
deleted file mode 100644
index e6bb8bc..0000000
--- a/ATRI/plugins/bilibili_dynamic/db.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from ATRI.database import BilibiliSubscription
-
-
-class DB:
- async def __aenter__(self):
- return self
-
- async def __aexit__(self, exc_type, exc_val, exc_tb):
- pass
-
- async def add_sub(self, uid: int, group_id: int):
- await BilibiliSubscription.create(uid=uid, group_id=group_id)
-
- async def update_sub(self, uid: int, group_id: int, update_map: dict):
- await BilibiliSubscription.filter(uid=uid, group_id=group_id).update(
- **update_map
- )
-
- async def del_sub(self, query_map: dict):
- await BilibiliSubscription.filter(**query_map).delete()
-
- async def get_sub_list(self, query_map: dict) -> list:
- return await BilibiliSubscription.filter(**query_map)
-
- async def get_all_subs(self) -> list:
- return await BilibiliSubscription.all()
diff --git a/ATRI/plugins/rss/rss_mikanan/data_source.py b/ATRI/plugins/rss/rss_mikanan/data_source.py
index 6aec1b5..289d4f3 100644
--- a/ATRI/plugins/rss/rss_mikanan/data_source.py
+++ b/ATRI/plugins/rss/rss_mikanan/data_source.py
@@ -2,44 +2,40 @@ import xmltodict
from ATRI.exceptions import RssError
from ATRI.utils import request, gen_random_str
+from ATRI.database import DatabaseWrapper, RssMikananiSubcription
-from .db import DB
+DB = DatabaseWrapper(RssMikananiSubcription)
class RssMikananSubscriptor:
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.mikan: 添加订阅失败")
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.mikan: 更新订阅失败")
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.mikan: 删除订阅失败")
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.mikan: 获取订阅列表失败")
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.mikan: 获取所有订阅失败")
diff --git a/ATRI/plugins/rss/rss_mikanan/db.py b/ATRI/plugins/rss/rss_mikanan/db.py
deleted file mode 100644
index ac3385d..0000000
--- a/ATRI/plugins/rss/rss_mikanan/db.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from ATRI.database import RssMikananiSubcription
-
-
-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 RssMikananiSubcription.create(_id=_id, group_id=group_id)
-
- async def update_sub(self, _id: str, group_id: int, update_map: dict):
- await RssMikananiSubcription.filter(_id=_id, group_id=group_id).update(
- **update_map
- )
-
- async def del_sub(self, query_map: dict):
- await RssMikananiSubcription.filter(**query_map).delete()
-
- async def get_sub_list(self, query_map: dict) -> list:
- return await RssMikananiSubcription.filter(**query_map)
-
- async def get_all_subs(self) -> list:
- return await RssMikananiSubcription.all()
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()
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()
Powered by cgit, Theme cgit-dark created by Ted Yin, updated by Mole Shang.