meida_backend/packages/kb/repo.py

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