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