导入导出 cookie 优化代码

This commit is contained in:
suyiiyii 2024-09-23 13:49:39 +08:00
parent 5d160dcadc
commit b04cbc2ced
3 changed files with 15 additions and 21 deletions

View File

@ -111,14 +111,7 @@ async def magic_cookie_gen(nbesf_data: SubGroup):
logger.info("开始添加 Cookie 流程")
for cookie in nbesf_data.cookies:
try:
new_cookie = DBCookie(
site_name=cookie.site_name,
content=cookie.content,
cookie_name=cookie.cookie_name,
cd_milliseconds=cookie.cd_milliseconds,
is_universal=cookie.is_universal,
tags=cookie.tags,
)
new_cookie = DBCookie(**model_dump(cookie, exclude={"targets"}))
cookie_id = await config.add_cookie(new_cookie)
for target in cookie.targets:
await config.add_cookie_target(target.target, target.platform_name, cookie_id)

View File

@ -10,7 +10,7 @@ from nonebot.compat import type_validate_python
from nonebot_plugin_datastore.db import create_session
from sqlalchemy.orm.strategy_options import selectinload
from .utils import NBESFVerMatchErr
from .utils import NBESFVerMatchErr, row2dict
from .nbesf_model import NBESFBase, v1, v2, v3
from ..db_model import User, Cookie, Target, Subscribe, CookieTarget
@ -59,7 +59,7 @@ async def subscribes_export(selector: Callable[[Select], Select]) -> v3.SubGroup
)
cookie_target_data = await sess.scalars(cookie_target_stmt)
cookie_target_dict = defaultdict(list)
cookie_target_dict: dict[Cookie, list[v3.Target]] = defaultdict(list)
for cookie_target in cookie_target_data:
target_payload = type_validate_python(v3.Target, cookie_target.target)
cookie_target_dict[cookie_target.cookie].append(target_payload)
@ -67,17 +67,10 @@ async def subscribes_export(selector: Callable[[Select], Select]) -> v3.SubGroup
cookies: list[v3.Cookie] = []
for cookie, targets in cookie_target_dict.items():
assert isinstance(cookie, Cookie)
cookies.append(
v3.Cookie(
site_name=cookie.site_name,
content=cookie.content,
cookie_name=cookie.cookie_name,
cd_milliseconds=cookie.cd_milliseconds,
is_universal=cookie.is_universal,
tags=cookie.tags,
targets=targets,
)
)
cookie_dict = row2dict(cookie)
cookie_dict["tags"] = cookie.tags
cookie_dict["targets"] = targets
cookies.append(v3.Cookie(**cookie_dict))
sub_group = v3.SubGroup(groups=groups, cookies=cookies)

View File

@ -2,3 +2,11 @@ class NBESFVerMatchErr(Exception): ...
class NBESFParseErr(Exception): ...
def row2dict(row):
d = {}
for column in row.__table__.columns:
d[column.name] = str(getattr(row, column.name))
return d