suyiiyii/nonebot-bison
1
0
Fork 0
mirror of https://github.com/suyiiyii/nonebot-bison.git synced 2026-05-11 03:18:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

♻️ 重构 get_cookie 方法

Browse Source
This commit is contained in:
suyiiyii
2024-09-13 11:34:45 +08:00
parent 4b8d6a9379
commit 4791fb69e0
6 changed files with 73 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nonebot_bison/config/db_config.py
+16 -35
View File
@@ -259,36 +259,32 @@ class DBConfig:
)
return res
async def get_cookie(self, site_name: str = None, target: T_Target = None) -> Sequence[Cookie]:
"""根据 site_name 和 target 获取 cookie，不会返回匿名cookie"""
async def get_cookie(
self,
site_name: str | None = None,
target: T_Target | None = None,
is_universal: bool | None = None,
is_anonymous: bool | None = None,
) -> Sequence[Cookie]:
"""获取满足传入条件的所有 cookie"""
async with create_session() as sess:
query = select(Cookie).distinct().where(Cookie.is_universal == False) # noqa: E712
query = select(Cookie).distinct()
if is_universal is not None:
query = query.where(Cookie.is_universal == is_universal)
if is_anonymous is not None:
query = query.where(Cookie.is_anonymous == is_anonymous)
if site_name:
query = query.where(Cookie.site_name == site_name)
query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets))
res = (await sess.scalars(query)).all()
if target:
# 如果指定了 target，过滤掉不满足要求的cookie
query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target)
ids = set((await sess.scalars(query)).all())
res = [cookie for cookie in res if cookie.id in ids]
# 如果指定了 target 且未指定 is_universal，则添加返回 universal cookie
res = [cookie for cookie in res if cookie.id in ids or cookie.is_universal]
return res
async def get_unviersal_cookie(self, site_name: str = None) -> Sequence[Cookie]:
async with create_session() as sess:
query = select(Cookie).distinct().where(Cookie.is_universal == True) # noqa: E712
if site_name:
query = query.where(Cookie.site_name == site_name)
res = (await sess.scalars(query)).all()
return res
async def add_cookie_with_content(self, site_name: str, content: str) -> int:
async with create_session() as sess:
cookie = Cookie(site_name=site_name, content=content)
sess.add(cookie)
await sess.commit()
await sess.refresh(cookie)
return cookie.id
async def add_cookie(self, cookie: Cookie) -> int:
async with create_session() as sess:
sess.add(cookie)
@@ -320,21 +316,6 @@ class DBConfig:
await sess.execute(delete(Cookie).where(Cookie.id == cookie_id))
await sess.commit()
async def get_cookie_by_target(self, target: T_Target, site_name: str) -> Sequence[Cookie]:
async with create_session() as sess:
query = (
select(Cookie)
.join(CookieTarget)
.join(Target)
.where(Target.site_name == site_name, Target.target == target)
)
return (await sess.scalars(query)).all()
async def get_universal_cookie(self, site_name: str) -> Sequence[Cookie]:
async with create_session() as sess:
query = select(Cookie).where(Cookie.site_name == site_name).where(Cookie.is_universal == True) # noqa: E712
return (await sess.scalars(query)).all()
async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
"""通过 cookie_id 可以唯一确定一个 Cookie，通过 target 和 platform_name 可以唯一确定一个 Target"""
async with create_session() as sess: