🚚💩 更改数据库位置 & 优化代码

10 files changed, 86 insertions, 70 deletions

diff --git a/ATRI/__init__.py b/ATRI/__init__.py
index 7f17778..fd166fe 100644
--- a/ATRI/__init__.py
+++ b/ATRI/__init__.py
@@ -4,6 +4,7 @@ import nonebot
 from nonebot.adapters.onebot.v11 import Adapter
 
 from .config import RUNTIME_CONFIG, InlineGoCQHTTP
+from .database import init_database
 
 __version__ = "YHN-001-A05.fix1"
 
@@ -22,6 +23,7 @@ def init():
     nonebot.load_plugins("ATRI/plugins")
     if InlineGoCQHTTP.enabled:
         nonebot.load_plugin("nonebot_plugin_gocqhttp")
+    init_database()
     sleep(3)
 
 
diff --git a/ATRI/database/__init__.py b/ATRI/database/__init__.py
new file mode 100644
index 0000000..6c58d20
--- /dev/null
+++ b/ATRI/database/__init__.py
@@ -0,0 +1,2 @@
+from .db import init_database
+from .models import BilibiliSubscription, TwitterSubscription
diff --git a/ATRI/database/db.py b/ATRI/database/db.py
new file mode 100644
index 0000000..6e20ad3
--- /dev/null
+++ b/ATRI/database/db.py
@@ -0,0 +1,27 @@
+from pathlib import Path
+from tortoise import Tortoise, run_async
+
+from ATRI.log import logger as log
+
+
+# 临时的实现，寻求更好的方式！欢迎pr
+
+
+DB_DIR = Path(".") / "data" / "sql"
+DB_DIR.mkdir(parents=True, exist_ok=True)
+
+
+async def run():
+    from ATRI.database import models
+
+    await Tortoise.init(
+        db_url=f"sqlite://{DB_DIR}/db.sqlite3",
+        modules={"models": [locals()["models"]]},
+    )
+    await Tortoise.generate_schemas()
+
+
+def init_database():
+    log.info("正在初始化数据库...")
+    run_async(run())
+    log.success("数据库初始化完成")
diff --git a/ATRI/database/models.py b/ATRI/database/models.py
new file mode 100644
index 0000000..e268f01
--- /dev/null
+++ b/ATRI/database/models.py
@@ -0,0 +1,19 @@
+from datetime import datetime
+
+from tortoise import fields
+from tortoise.models import Model
+
+
+class BilibiliSubscription(Model):
+    uid = fields.IntField()
+    group_id = fields.IntField()
+    up_nickname = fields.TextField(null=True)
+    last_update = fields.DatetimeField(default=datetime.fromordinal(1))
+
+
+class TwitterSubscription(Model):
+    tid = fields.IntField()
+    group_id = fields.IntField()
+    name = fields.TextField(null=True)
+    screen_name = fields.TextField(null=True)
+    last_update = fields.DatetimeField(default=datetime.fromordinal(1))
diff --git a/ATRI/plugins/bilibili_dynamic/__init__.py b/ATRI/plugins/bilibili_dynamic/__init__.py
index fb3b7eb..cd2f0d5 100644
--- a/ATRI/plugins/bilibili_dynamic/__init__.py
+++ b/ATRI/plugins/bilibili_dynamic/__init__.py
@@ -16,9 +16,9 @@ from nonebot.adapters.onebot.v11 import Message, GroupMessageEvent
 from ATRI.log import logger as log
 from ATRI.utils import timestamp2datetime
 from ATRI.utils.apscheduler import scheduler
+from ATRI.database import BilibiliSubscription
 
 from .data_source import BilibiliDynamicSubscriptor
-from .database.models import Subscription
 
 
 add_sub = BilibiliDynamicSubscriptor().cmd_as_group("add", "添加b站up主订阅")
@@ -121,7 +121,7 @@ tq = asyncio.Queue()
 class BilibiliDynamicChecker(BaseTrigger):
     def get_next_fire_time(self, previous_fire_time, now):
         sub = BilibiliDynamicSubscriptor()
-        conf = sub.load_service("b站动态订阅-rebu")
+        conf = sub.load_service("b站动态订阅")
         if conf.get("enabled"):
             return now
 
@@ -132,14 +132,19 @@ class BilibiliDynamicChecker(BaseTrigger):
     max_instances=3,  # type: ignore
     misfire_grace_time=60,  # type: ignore
 )
-async def _check_dynamic():
+async def _check_bd():
     sub = BilibiliDynamicSubscriptor()
-    all_dy = await sub.get_all_subs()
+    try:
+        all_dy = await sub.get_all_subs()
+    except Exception:
+        log.debug("b站订阅列表为空 跳过")
+        return
+    
     if tq.empty():
         for i in all_dy:
             await tq.put(i)
     else:
-        m: Subscription = tq.get_nowait()
+        m: BilibiliSubscription = tq.get_nowait()
         log.info(f"准备查询up主[{m.up_nickname}]的动态，队列剩余 {tq.qsize()}")
 
         ts = int(m.last_update.timestamp())
@@ -149,6 +154,7 @@ async def _check_dynamic():
             result = sub.extract_dyanmic(info["cards"])
         if not result:
             log.warning(f"无法获取up主[{m.up_nickname}]的动态")
+            return
 
         for i in result:
             i["name"] = m.up_nickname
diff --git a/ATRI/plugins/bilibili_dynamic/data_source.py b/ATRI/plugins/bilibili_dynamic/data_source.py
index 28816bb..cd6db37 100644
--- a/ATRI/plugins/bilibili_dynamic/data_source.py
+++ b/ATRI/plugins/bilibili_dynamic/data_source.py
@@ -9,7 +9,7 @@ from ATRI.rule import is_in_service
 from ATRI.utils import timestamp2datetime
 from ATRI.exceptions import BilibiliDynamicError
 
-from .database import DB
+from .db import DB
 from .api import API
 
 
diff --git a/ATRI/plugins/bilibili_dynamic/database/__init__.py b/ATRI/plugins/bilibili_dynamic/database/__init__.py
deleted file mode 100644
index 6840881..0000000
--- a/ATRI/plugins/bilibili_dynamic/database/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .db import DB
diff --git a/ATRI/plugins/bilibili_dynamic/database/db.py b/ATRI/plugins/bilibili_dynamic/database/db.py
deleted file mode 100644
index 37ed223..0000000
--- a/ATRI/plugins/bilibili_dynamic/database/db.py
+++ /dev/null
@@ -1,49 +0,0 @@
-from pathlib import Path
-from tortoise import Tortoise
-
-from ATRI import driver
-from .models import Subscription
-
-
-DB_DIR = Path(".") / "data" / "database" / "bilibili_dynamic"
-DB_DIR.mkdir(parents=True, exist_ok=True)
-
-
-class DB:
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        pass
-
-    async def init(self):
-        from . import models
-
-        await Tortoise.init(
-            db_url=f"sqlite://{DB_DIR}/db.sqlite3",
-            modules={"models": [locals()["models"]]},
-        )
-        await Tortoise.generate_schemas()
-
-    async def add_sub(self, uid: int, group_id: int):
-        await Subscription.create(uid=uid, group_id=group_id)
-
-    async def update_sub(self, uid: int, update_map: dict):
-        await Subscription.filter(uid=uid).update(**update_map)
-
-    async def del_sub(self, query_map: dict):
-        await Subscription.filter(**query_map).delete()
-
-    async def get_sub_list(self, query_map: dict) -> list:
-        return await Subscription.filter(**query_map)
-
-    async def get_all_subs(self) -> list:
-        return await Subscription.all()
-
-
-async def init():
-    async with DB() as db:
-        await db.init()
-
-
-driver().on_startup(init)
diff --git a/ATRI/plugins/bilibili_dynamic/database/models.py b/ATRI/plugins/bilibili_dynamic/database/models.py
deleted file mode 100644
index d0cb3dd..0000000
--- a/ATRI/plugins/bilibili_dynamic/database/models.py
+++ /dev/null
@@ -1,14 +0,0 @@
-from datetime import datetime
-
-from tortoise import fields
-from tortoise.models import Model
-
-
-class Subscription(Model):
-    uid = fields.IntField()
-    group_id = fields.IntField()
-    up_nickname = fields.TextField(null=True)
-    last_update = fields.DatetimeField(default=datetime.fromordinal(1))
-
-    def __str__(self) -> str:
-        return f"[{self.uid}|{self.group_id}|{self.up_nickname}|{self.last_update}]"
diff --git a/ATRI/plugins/bilibili_dynamic/db.py b/ATRI/plugins/bilibili_dynamic/db.py
new file mode 100644
index 0000000..a969177
--- /dev/null
+++ b/ATRI/plugins/bilibili_dynamic/db.py
@@ -0,0 +1,24 @@
+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, update_map: dict):
+        await BilibiliSubscription.filter(uid=uid).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()