xsl
9c1802130d
第三批（T12 后台/T13 黄金集/T14 检索缓存）协调提交： - kb：KbAssetService（登记/审核，tryon_ready 只能随 approved 打开）+ Asset 模型与 repo - retrieval：list_gaps 缺口看板查询 + invalidate_cache() 钩子（T14 背后实现真缓存） - 配置：ADMIN_TOKEN（空=后台 404）+ 检索缓存三参数 - pytest golden marker（默认跳过，CI 强制跑）+ Makefile golden + CI 防线 5c - 测试隔离修复：test_auth_sms 加 autouse 清 sms:* 键夹具（频控计数跨运行 累积会毒化后续运行——随机 IP 空间仅 254 且部分用例固定 device_id）
205 lines
7.3 KiB
Python
205 lines
7.3 KiB
Python
"""kb 域数据访问层。**只有本模块允许触碰 kb_entity / kb_relation / kb_color。**
|
||
|
||
``traverse`` 用递归 CTE 而不是在 Python 里逐跳发查询：TD-09 的结论是
|
||
知识图谱的查询都是「从某实体出发的浅层有界遍历」（1~3 跳），邻接表 +
|
||
CTE 在 PG 内一次往返就能完成，不值得为此引入图数据库的第二套存储与运维。
|
||
防环在 CTE 里做（``path`` 数组 + ``@>`` 包含判断），保证有环数据
|
||
（A→B→A）也不会让递归发散。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Iterable
|
||
from dataclasses import dataclass
|
||
|
||
from sqlalchemy import ColumnElement, func, select, text
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from packages.shared.models import KbAsset, KbColor, KbEntity, KbRelation
|
||
|
||
#: 递归遍历 SQL（1~3 跳）。``path`` 记录已经过的实体 id，``rels`` 与逐跳
|
||
#: rel_type 平行——每行结果本身就是一条完整可解释路径，不需要在 Python
|
||
#: 侧再拼接。深度上限由参数 ``:max_hops`` 控制（service 层已校验 1~3）。
|
||
#: ``::varchar[]`` 显转型不可省：锚点项推得 varchar(64)[]、递归项推得
|
||
#: varchar[]，PG 要求两段类型完全一致，不转型直接 DatatypeMismatchError。
|
||
_TRAVERSE_SQL = text(
|
||
"""
|
||
WITH RECURSIVE walk(from_id, to_id, rel_type, hop, path, rels) AS (
|
||
SELECT r.from_id, r.to_id, r.rel_type, 1,
|
||
ARRAY[r.from_id, r.to_id]::bigint[], ARRAY[r.rel_type]::varchar[]
|
||
FROM kb_relation r
|
||
WHERE r.from_id = :start_id
|
||
UNION ALL
|
||
SELECT r.from_id, r.to_id, r.rel_type, w.hop + 1,
|
||
w.path || r.to_id, w.rels || ARRAY[r.rel_type]::varchar[]
|
||
FROM kb_relation r
|
||
JOIN walk w ON r.from_id = w.to_id
|
||
WHERE w.hop < :max_hops
|
||
AND NOT w.path @> ARRAY[r.to_id]
|
||
)
|
||
SELECT hop, path, rels FROM walk ORDER BY hop, path, rels
|
||
"""
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class RawPath:
|
||
"""CTE 直接吐出的一条路径（防环已由 SQL 保证，len(rel_types)==len(node_ids)-1）。"""
|
||
|
||
node_ids: tuple[int, ...]
|
||
rel_types: tuple[str, ...]
|
||
|
||
|
||
# --- kb_entity ---------------------------------------------------------------
|
||
|
||
|
||
async def get_entity(session: AsyncSession, entity_id: int) -> KbEntity | None:
|
||
return await session.get(KbEntity, entity_id)
|
||
|
||
|
||
async def get_entity_for_update(session: AsyncSession, entity_id: int) -> KbEntity | None:
|
||
result = await session.execute(
|
||
select(KbEntity).where(KbEntity.entity_id == entity_id).with_for_update()
|
||
)
|
||
return result.scalar_one_or_none()
|
||
|
||
|
||
async def get_entities_by_ids(session: AsyncSession, ids: Iterable[int]) -> dict[int, KbEntity]:
|
||
id_list = list(ids)
|
||
if not id_list:
|
||
return {}
|
||
result = await session.execute(select(KbEntity).where(KbEntity.entity_id.in_(id_list)))
|
||
return {entity.entity_id: entity for entity in result.scalars()}
|
||
|
||
|
||
async def find_entity_by_name(
|
||
session: AsyncSession, *, entity_type: str, name: str
|
||
) -> KbEntity | None:
|
||
# UNIQUE(entity_type, name) 保证至多一行
|
||
result = await session.execute(
|
||
select(KbEntity).where(KbEntity.entity_type == entity_type, KbEntity.name == name)
|
||
)
|
||
return result.scalar_one_or_none()
|
||
|
||
|
||
def add_entity(session: AsyncSession, entity: KbEntity) -> KbEntity:
|
||
session.add(entity)
|
||
return entity
|
||
|
||
|
||
async def search_entities(
|
||
session: AsyncSession,
|
||
*,
|
||
entity_type: str | None = None,
|
||
status: str | None = None,
|
||
name_ilike: str | None = None,
|
||
limit: int = 20,
|
||
offset: int = 0,
|
||
) -> tuple[list[KbEntity], int]:
|
||
"""过滤 + 分页。返回 ``(当页行, 满足过滤条件的总数)``。
|
||
|
||
``name_ilike`` 不转义 ``%``/``_``：这是后台运营的模糊搜索框，
|
||
通配符被当通配符用符合直觉，也不是注入面（参数化绑定）。
|
||
"""
|
||
conds: list[ColumnElement[bool]] = []
|
||
if entity_type is not None:
|
||
conds.append(KbEntity.entity_type == entity_type)
|
||
if status is not None:
|
||
conds.append(KbEntity.status == status)
|
||
if name_ilike:
|
||
conds.append(KbEntity.name.ilike(f"%{name_ilike}%"))
|
||
|
||
total = int(await session.scalar(select(func.count()).select_from(KbEntity).where(*conds)) or 0)
|
||
result = await session.execute(
|
||
select(KbEntity)
|
||
.where(*conds)
|
||
.order_by(KbEntity.created_at.desc(), KbEntity.entity_id.desc())
|
||
.limit(limit)
|
||
.offset(offset)
|
||
)
|
||
return list(result.scalars().all()), total
|
||
|
||
|
||
# --- kb_relation -------------------------------------------------------------
|
||
|
||
|
||
async def find_relation(
|
||
session: AsyncSession, *, from_id: int, rel_type: str, to_id: int
|
||
) -> KbRelation | None:
|
||
result = await session.execute(
|
||
select(KbRelation).where(
|
||
KbRelation.from_id == from_id,
|
||
KbRelation.rel_type == rel_type,
|
||
KbRelation.to_id == to_id,
|
||
)
|
||
)
|
||
return result.scalar_one_or_none()
|
||
|
||
|
||
def add_relation(session: AsyncSession, relation: KbRelation) -> KbRelation:
|
||
session.add(relation)
|
||
return relation
|
||
|
||
|
||
async def traverse(session: AsyncSession, *, from_id: int, max_hops: int) -> list[RawPath]:
|
||
"""从 ``from_id`` 出发沿 from→to 方向做 ``max_hops`` 跳的有界遍历。
|
||
|
||
深度合法性（1~3）由 service 层校验，这里信任调用方。返回按
|
||
（跳数, 路径, 关系序列）排序的全部路径——同一终点多条路径都保留，
|
||
「为什么推荐」恰恰需要把每条依据都摆出来。
|
||
"""
|
||
result = await session.execute(_TRAVERSE_SQL, {"start_id": from_id, "max_hops": max_hops})
|
||
return [RawPath(node_ids=tuple(row.path), rel_types=tuple(row.rels)) for row in result]
|
||
|
||
|
||
# --- kb_color ----------------------------------------------------------------
|
||
|
||
|
||
async def list_colors(session: AsyncSession, *, season_code: str | None = None) -> list[KbColor]:
|
||
stmt = select(KbColor).order_by(KbColor.season_code, KbColor.color_id)
|
||
if season_code is not None:
|
||
stmt = stmt.where(KbColor.season_code == season_code)
|
||
result = await session.execute(stmt)
|
||
return list(result.scalars().all())
|
||
|
||
|
||
def add_color(session: AsyncSession, color: KbColor) -> KbColor:
|
||
session.add(color)
|
||
return color
|
||
|
||
|
||
async def get_asset(session: AsyncSession, asset_id: int) -> KbAsset | None:
|
||
return await session.get(KbAsset, asset_id)
|
||
|
||
|
||
async def list_assets(
|
||
session: AsyncSession,
|
||
*,
|
||
entity_id: int | None = None,
|
||
tryon_ready: bool | None = None,
|
||
limit: int = 50,
|
||
offset: int = 0,
|
||
) -> tuple[list[KbAsset], int]:
|
||
"""素材列表（含总数）。tryon_ready 过滤走 partial 索引。"""
|
||
from sqlalchemy import func, select
|
||
|
||
conditions = []
|
||
if entity_id is not None:
|
||
conditions.append(KbAsset.entity_id == entity_id)
|
||
if tryon_ready is not None:
|
||
conditions.append(KbAsset.tryon_ready == tryon_ready)
|
||
|
||
base = select(KbAsset)
|
||
if conditions:
|
||
base = base.where(*conditions)
|
||
total = (await session.execute(select(func.count()).select_from(base.subquery()))).scalar_one()
|
||
rows = (
|
||
await session.execute(base.order_by(KbAsset.asset_id.desc()).limit(limit).offset(offset))
|
||
).scalars()
|
||
return list(rows), int(total)
|
||
|
||
|
||
def add_asset(session: AsyncSession, asset: KbAsset) -> KbAsset:
|
||
session.add(asset)
|
||
return asset
|