content-find-agent-new/content_agent/integrations/pattern_pg.py

"""M9C:只读直连 pattern PG(open_aigc),判需求是否含分类树终端元素(叶子)。

叶子 = pattern_mining_category 中某 id 没有任何行 parent_id=该 id(同 execution)。
Gate 1:需求 itemset_items 的 category_id 里至少有一个叶子 → 放行;否则不做。
"""

from __future__ import annotations

from typing import Any

from content_agent.errors import ContentAgentError, ErrorCode
from content_agent.integrations import timeout_config
from content_agent.integrations.crawapi_http import _env, _load_env_file

_LEAF_SQL = (
    "SELECT 1 FROM pattern_mining_category c "
    "WHERE c.execution_id = %s AND c.id = ANY(%s) "
    "AND NOT EXISTS (SELECT 1 FROM pattern_mining_category ch "
    "WHERE ch.execution_id = c.execution_id AND ch.parent_id = c.id) LIMIT 1"
)


class PatternPgClient:
    def __init__(
        self,
        *,
        host: str,
        port: int,
        user: str,
        password: str,
        database: str,
        timeout_seconds: float = 30.0,
    ) -> None:
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.timeout_seconds = timeout_seconds

    @classmethod
    def from_env(cls, env_path: str = ".env") -> "PatternPgClient":
        env = _load_env_file(env_path)
        return cls(
            host=_env("OPEN_AIGC_PG_HOST", env, required=True),
            port=int(_env("OPEN_AIGC_PG_PORT", env, default="5432")),
            user=_env("OPEN_AIGC_PG_USER", env, required=True),
            password=_env("OPEN_AIGC_PG_PASSWORD", env, required=True),
            database=_env("OPEN_AIGC_PG_DB_NAME", env, default="open_aigc"),
            timeout_seconds=timeout_config.total_timeout("pg", env=env),
        )

    def has_terminal_element(self, execution_id: int, category_ids: list[int]) -> bool:
        ids = [int(c) for c in category_ids if c is not None]
        if not ids:
            return True  # 无可判定 category → 不阻断(交给后续判定)
        import pg8000.dbapi  # 惰性导入:mock 测试用 fake client,不触发依赖

        try:
            conn = pg8000.dbapi.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                ssl_context=None,  # 服务器拒绝 SSL,必须明确关闭
                timeout=self.timeout_seconds,
            )
        except Exception as exc:  # PG 不可达:明确报错,不静默放过
            raise ContentAgentError(
                ErrorCode.DB_CONNECTION_FAILED,
                "pattern pg unreachable for gate1 terminal-element check",
                {"exception_type": type(exc).__name__},
            ) from exc
        try:
            cur = conn.cursor()
            # connect timeout 只管握手;execute/fetch 用服务端 statement_timeout 兜住,防查询永久阻塞。
            # pg8000/PG 不支持在 SET statement_timeout 里用 $1 占位符;这里先转 int 再内联。
            timeout_ms = max(1, int(self.timeout_seconds * 1000))
            cur.execute(f"SET statement_timeout = {timeout_ms}")
            cur.execute(_LEAF_SQL, (int(execution_id), ids))
            return cur.fetchone() is not None
        except Exception as exc:
            raise ContentAgentError(
                ErrorCode.DB_SCHEMA_NOT_READY,
                "pattern pg gate1 terminal-element query failed",
                {
                    "exception_type": type(exc).__name__,
                    "error_message": str(exc)[:500],
                    "execution_id": int(execution_id),
                    "category_id_count": len(ids),
                },
            ) from exc
        finally:
            conn.close()