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

3 files changed, 48 insertions, 0 deletions

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))