huoyan-enterprise/backend/core/graph_metadata.py

"""
图谱元数据表名（graphs / star_graph）与 chat_threads 知识图谱外键列名兼容。
"""
from __future__ import annotations

import asyncio
from typing import Final, Optional

import asyncpg

from logger.logging import get_logger

logger = get_logger(__name__)

_ALLOWED_TABLES: Final[frozenset[str]] = frozenset({"graphs", "star_graph"})
_ALLOWED_KG_COLS: Final[frozenset[str]] = frozenset({"knowledge_graph_id", "novel_graph_id"})

_lock = asyncio.Lock()
_ready: bool = False

GRAPH_TABLE: str = "graphs"
# None = 未探测（或库表无知识图谱外键列）；ensure_graph_metadata 后会设为实际列名或保持 None
CHAT_THREAD_KG_COLUMN: Optional[str] = None
# chat_threads 是否存在 ip 列（应用 INSERT 会话时写入；无此列则省略，避免 INSERT 失败导致会话列表永远为空）
CHAT_THREADS_HAS_IP_COLUMN: bool = False
# chat_threads 是否存在 llm_provider / llm_model（记录会话最近选用模型；见 migrations/add_chat_threads_llm_columns.sql）
CHAT_THREADS_HAS_LLM_COLUMNS: bool = False


def graph_table_sql() -> str:
    if GRAPH_TABLE not in _ALLOWED_TABLES:
        raise RuntimeError(f"invalid GRAPH_TABLE: {GRAPH_TABLE!r}")
    return GRAPH_TABLE


def chat_thread_kg_column_sql() -> str:
    """返回 chat_threads 上绑定图谱的列名；若库中无该列则抛错（仅用于确需写入该列的路径）。"""
    if CHAT_THREAD_KG_COLUMN is None:
        raise RuntimeError(
            "chat_threads 缺少 knowledge_graph_id / novel_graph_id 列，请执行 migrations/knowledge_graph_and_processing.sql"
        )
    if CHAT_THREAD_KG_COLUMN not in _ALLOWED_KG_COLS:
        raise RuntimeError(f"invalid CHAT_THREAD_KG_COLUMN: {CHAT_THREAD_KG_COLUMN!r}")
    return CHAT_THREAD_KG_COLUMN


def chat_thread_kg_select_fragment_sql() -> str:
    """用于 SELECT 列表：无物理列时返回 NULL，避免引用不存在的列导致会话列表等接口 500。"""
    if CHAT_THREAD_KG_COLUMN is None:
        return "NULL::integer AS knowledge_graph_id"
    if CHAT_THREAD_KG_COLUMN not in _ALLOWED_KG_COLS:
        raise RuntimeError(f"invalid CHAT_THREAD_KG_COLUMN: {CHAT_THREAD_KG_COLUMN!r}")
    return f"{CHAT_THREAD_KG_COLUMN} AS knowledge_graph_id"


def chat_threads_has_kg_column() -> bool:
    return CHAT_THREAD_KG_COLUMN is not None


def chat_threads_has_ip_column() -> bool:
    return CHAT_THREADS_HAS_IP_COLUMN


def chat_threads_has_llm_columns() -> bool:
    return CHAT_THREADS_HAS_LLM_COLUMNS


def chat_thread_llm_select_fragment_sql() -> str:
    """SELECT 列表片段：无列时返回 NULL，避免未迁移库 500。"""
    if CHAT_THREADS_HAS_LLM_COLUMNS:
        return "llm_provider, llm_model"
    return "NULL::varchar AS llm_provider, NULL::varchar AS llm_model"


async def ensure_graph_metadata(conn: asyncpg.Connection) -> None:
    """首次连接数据库时解析表名与 chat_threads 列名（仅白名单）。"""
    global _ready, GRAPH_TABLE, CHAT_THREAD_KG_COLUMN, CHAT_THREADS_HAS_IP_COLUMN, CHAT_THREADS_HAS_LLM_COLUMNS
    if _ready:
        return
    async with _lock:
        if _ready:
            return
        has_g = await conn.fetchval(
            """
            SELECT EXISTS (
              SELECT 1 FROM information_schema.tables
              WHERE table_schema = 'public' AND table_name = 'graphs'
            )
            """
        )
        has_s = await conn.fetchval(
            """
            SELECT EXISTS (
              SELECT 1 FROM information_schema.tables
              WHERE table_schema = 'public' AND table_name = 'star_graph'
            )
            """
        )
        if has_g:
            GRAPH_TABLE = "graphs"
        elif has_s:
            GRAPH_TABLE = "star_graph"
            logger.info("图谱元数据表使用 PostgreSQL 表名 star_graph，建议统一为 graphs")
        else:
            GRAPH_TABLE = "graphs"
            logger.warning("未找到 public.graphs 或 public.star_graph，请先执行数据库迁移")

        has_kg = await conn.fetchval(
            """
            SELECT EXISTS (
              SELECT 1 FROM information_schema.columns
              WHERE table_schema = 'public' AND table_name = 'chat_threads'
                AND column_name = 'knowledge_graph_id'
            )
            """
        )
        has_ng = await conn.fetchval(
            """
            SELECT EXISTS (
              SELECT 1 FROM information_schema.columns
              WHERE table_schema = 'public' AND table_name = 'chat_threads'
                AND column_name = 'novel_graph_id'
            )
            """
        )
        if has_kg:
            CHAT_THREAD_KG_COLUMN = "knowledge_graph_id"
        elif has_ng:
            CHAT_THREAD_KG_COLUMN = "novel_graph_id"
            logger.info("chat_threads 使用列 novel_graph_id，可迁移为 knowledge_graph_id")
        else:
            CHAT_THREAD_KG_COLUMN = None
            logger.warning(
                "chat_threads 未找到 knowledge_graph_id / novel_graph_id；会话列表仍可查询，图谱绑定需执行迁移"
            )

        _has_ip = await conn.fetchval(
            """
            SELECT EXISTS (
              SELECT 1 FROM information_schema.columns
              WHERE table_schema = 'public' AND table_name = 'chat_threads'
                AND column_name = 'ip'
            )
            """
        )
        CHAT_THREADS_HAS_IP_COLUMN = bool(_has_ip)

        _has_llm = await conn.fetchval(
            """
            SELECT EXISTS (
              SELECT 1 FROM information_schema.columns
              WHERE table_schema = 'public' AND table_name = 'chat_threads'
                AND column_name = 'llm_provider'
            )
            """
        )
        CHAT_THREADS_HAS_LLM_COLUMNS = bool(_has_llm)

        _ready = True


def reset_graph_metadata() -> None:
    global _ready, GRAPH_TABLE, CHAT_THREAD_KG_COLUMN, CHAT_THREADS_HAS_IP_COLUMN, CHAT_THREADS_HAS_LLM_COLUMNS
    _ready = False
    GRAPH_TABLE = "graphs"
    CHAT_THREAD_KG_COLUMN = None
    CHAT_THREADS_HAS_IP_COLUMN = False
    CHAT_THREADS_HAS_LLM_COLUMNS = False