✨ 通过 nb-cli 实现数据库一键导入导出 (#210)

* feat: 实现导出存储的订阅信息的功能

* test: 编写导出功能测试

* test: 使用tmp_path

* feat: 实现导入订阅文件功能

* refactor: 将订阅导入导出部分独立出来

* fix: 修复一些拼写错误
test: 完成import的第一个测试

* feat: 将订阅导入导出函数加入nb script

test: 添加cli测试

* test: 完善subs import测试

* 🐛 fix nb cli entrypoint name error

* fix: 修改错误的entry_point, 关闭yaml导出时对键名的排序

* fix: 使用更简短的命令名

* 🚚 将subs_io迁移到config下

* ♻️ 不再使用抛出异常的方式创建目录

* refactor: 将subscribe_export类转为函数

* refactor: 将subscribe_import类转为函数

* refactor: 根据重写的subs_io重新调整cli

* test: 调整重写subs_io后的test

* chore: 清理未使用的import内容

* feat(cli): 将--yaml更改为--format

* test: 调整测试

* fix(cli): 为import添加不支持格式的报错

* ⚡ improve export performace

* feat: subscribes_import函数不再需要传入参数函数，而是指定为add_subscribes

fix: nbesf_parser在传入str时将调用parse_raw, 否则调用parse_obj

* feat: subscribes_import现在会根据nbesf_data的版本选择合适的导入方式

* fix(test): 调整测试

* feat: nb bison export命令不再将文件导出到data目录，而是当前工作目录

* docs: 增添相关文档

* fix(test): 修复错误的变量名

---------

Co-authored-by: felinae98 <731499577@qq.com>

nonebot_bison/config/subs_io/__init__.py

@@ -0,0 +1,3 @@
+from .subs_io import nbesf_parser, subscribes_export, subscribes_import
+
+__all__ = ["subscribes_export", "subscribes_import", "nbesf_parser"]

nonebot_bison/config/subs_io/nbesf_model.py

@@ -0,0 +1,81 @@
+""" nbexf is Nonebot Bison Enchangable Subscribes File! """
+
+from pydantic import BaseModel
+
+from ...types import Category, Tag
+
+# ===== nbesf 定义格式 ====== #
+NBESF_VERSION = 1
+
+
+class UserHead(BaseModel, orm_mode=True):
+    """Bison快递包收货信息"""
+
+    type: str
+    uid: int
+
+    def __hash__(self):
+        return hash((self.type, self.uid))
+
+
+class Target(BaseModel, orm_mode=True):
+    """Bsion快递包发货信息"""
+
+    target_name: str
+    target: str
+    platform_name: str
+    default_schedule_weight: int
+
+
+class SubPayload(BaseModel, orm_mode=True):
+    """Bison快递包里的单件货物"""
+
+    categories: list[Category]
+    tags: list[Tag]
+    target: Target
+
+
+class SubPack(BaseModel):
+    """Bison给指定用户派送的快递包"""
+
+    user: UserHead
+    subs: list[SubPayload]
+
+
+class SubGroup(BaseModel):
+    """
+    Bison的全部订单(按用户分组)
+
+    结构参见`nbesf_model.py`
+    """
+
+    version: int = NBESF_VERSION  # 表示nbesf格式版本，从1开始
+    groups: list[SubPack]
+
+
+# ======================= #
+
+
+class SubReceipt(BaseModel):
+    """
+    快递包中每件货物的收据
+
+    导入订阅时的Model
+    """
+
+    user: int
+    user_type: str
+    target: str
+    target_name: str
+    platform_name: str
+    cats: list[Category]
+    tags: list[Tag]
+    # default_schedule_weight: int
+
+
+class NBESFVerMatchErr(Exception):
+    ...
+
+
+class NBESFParseErr(Exception):
+    ...

nonebot_bison/config/subs_io/subs_io.py

@@ -0,0 +1,88 @@
+from collections import defaultdict
+from functools import partial
+from typing import Any, Callable, TypeVar
+
+from nonebot.log import logger
+from nonebot_plugin_datastore.db import create_session
+from sqlalchemy import select
+from sqlalchemy.orm.strategy_options import selectinload
+from sqlalchemy.sql.selectable import Select
+
+from ..db_model import Subscribe, User
+from .nbesf_model import (
+    NBESFParseErr,
+    NBESFVerMatchErr,
+    SubGroup,
+    SubPack,
+    SubPayload,
+    UserHead,
+)
+from .utils import subs_receipt_gen_ver_1
+
+T = TypeVar("T", bound=Select)
+
+
+async def subscribes_export(selector: Callable[[T], T]) -> SubGroup:
+    """
+    将Bison订阅导出为 Nonebot Bison Exchangable Subscribes File 标准格式的 SubGroup 类型数据
+
+    selector:
+        对 sqlalchemy Select 对象的操作函数，用于限定查询范围 e.g. lambda stmt: stmt.where(User.uid=2233, User.type="group")
+    """
+    async with create_session() as sess:
+        sub_stmt = select(Subscribe).join(User)
+        sub_stmt = selector(sub_stmt).options(selectinload(Subscribe.target))
+        sub_data = await sess.scalars(sub_stmt)
+
+        user_stmt = select(User).join(Subscribe)
+        user_stmt = selector(user_stmt).distinct()
+        user_data = await sess.scalars(user_stmt)
+
+    groups: list[SubPack] = []
+    user_id_sub_dict: dict[int, list[SubPayload]] = defaultdict(list)
+
+    for sub in sub_data:
+        sub_paylaod = SubPayload.from_orm(sub)
+        user_id_sub_dict[sub.user_id].append(sub_paylaod)
+
+    for user in user_data:
+        user_head = UserHead.from_orm(user)
+        sub_pack = SubPack(user=user_head, subs=user_id_sub_dict[user.id])
+        groups.append(sub_pack)
+
+    sub_group = SubGroup(groups=groups)
+
+    return sub_group
+
+
+async def subscribes_import(
+    nbesf_data: SubGroup,
+):
+    """
+    从 Nonebot Bison Exchangable Subscribes File 标准格式的数据中导入订阅
+
+    nbesf_data:
+        符合nbesf_model标准的 SubGroup 类型数据
+    """
+
+    logger.info("开始添加订阅流程")
+    match nbesf_data.version:
+        case 1:
+            await subs_receipt_gen_ver_1(nbesf_data)
+        case _:
+            raise NBESFVerMatchErr(f"不支持的NBESF版本：{nbesf_data.version}")
+    logger.info("订阅流程结束，请检查所有订阅记录是否全部添加成功")
+
+
+def nbesf_parser(raw_data: Any) -> SubGroup:
+    try:
+        if isinstance(raw_data, str):
+            nbesf_data = SubGroup.parse_raw(raw_data)
+        else:
+            nbesf_data = SubGroup.parse_obj(raw_data)
+
+    except Exception as e:
+        logger.error("数据解析失败，该数据格式可能不满足NBESF格式标准！")
+        raise NBESFParseErr("数据解析失败") from e
+    else:
+        return nbesf_data

nonebot_bison/config/subs_io/utils.py

@@ -0,0 +1,28 @@
+from functools import partial
+
+from nonebot.log import logger
+
+from ..db_config import SubscribeDupException, config
+from .nbesf_model import NBESF_VERSION, SubGroup, SubReceipt
+
+
+async def subs_receipt_gen_ver_1(nbesf_data: SubGroup):
+    for item in nbesf_data.groups:
+        sub_receipt = partial(SubReceipt, user=item.user.uid, user_type=item.user.type)
+
+        for sub in item.subs:
+            receipt = sub_receipt(
+                target=sub.target.target,
+                target_name=sub.target.target_name,
+                platform_name=sub.target.platform_name,
+                cats=sub.categories,
+                tags=sub.tags,
+            )
+            try:
+                await config.add_subscribe(**receipt.dict())
+            except SubscribeDupException:
+                logger.warning(f"！添加订阅条目 {repr(receipt)} 失败: 相同的订阅已存在")
+            except Exception as e:
+                logger.error(f"！添加订阅条目 {repr(receipt)} 失败: {repr(e)}")
+            else:
+                logger.success(f"添加订阅条目 {repr(receipt)} 成功！")