Create-knowledge-find-decode/core/db.py

"""PG 连接基建：连 open_aigc（Greenplum / PG12 内核），供只读查询复用。

目前唯一用途是分类树 dump（scripts/dump_trees.py 查 public.global_category）。
连接 / 参数化风格对齐 ContentFindAgentNew 的 database_runtime.py。
"""
from __future__ import annotations

from typing import Any

import psycopg2
import psycopg2.extras

from core.config import PgConfig


def connect(cfg: PgConfig, schema: str | None = None) -> Any:
    """建连；schema 缺省用 cfg.schema，可显式传 public 查全局分类树。"""
    return psycopg2.connect(
        host=cfg.host,
        port=cfg.port,
        user=cfg.user,
        password=cfg.password,
        dbname=cfg.database,
        connect_timeout=cfg.timeout,
        options=f"-c search_path={schema or cfg.schema}",
    )


def fetch_all(cfg: PgConfig, sql: str, params: tuple = (), schema: str | None = None) -> list[dict]:
    """只读查询，返回 dict 列表。"""
    conn = connect(cfg, schema)
    try:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            return [dict(r) for r in cur.fetchall()]
    finally:
        conn.close()