Create-knowledge-find-decode/creation_knowledge/integrations/db.py

"""过程库 CkStore：写 / 读 creation_knowledge.ck_post、ck_knowledge_item。

目标实例是 Greenplum（PG12 内核），注意：
  - 不用 INSERT ... ON CONFLICT（Greenplum 分布表支持不稳）：upsert = 先 UPDATE，rowcount=0 再 INSERT。
  - 不更新分布键 id。
  - JSONB 列用 psycopg2.extras.Json 包装。
连接 / JSONB / 参数化风格对齐 ContentFindAgentNew 的 database_runtime.py。
"""
from __future__ import annotations

from typing import Any, Callable, Optional

import psycopg2
import psycopg2.extras

from creation_knowledge.config import PgConfig
from creation_knowledge.models import Post

ConnectionFactory = Callable[[], Any]


def _connect(cfg: PgConfig) -> Any:
    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={cfg.schema}",
    )


class CkStore:
    def __init__(
        self,
        config: PgConfig,
        connection_factory: ConnectionFactory | None = None,
    ) -> None:
        self.config = config
        self._connection_factory = connection_factory or (lambda: _connect(config))

    # ---------- 写 ----------
    def upsert_post(self, post: Post) -> None:
        """落 fetch 产物：raw 等基础字段，stage=fetched。先 UPDATE 后 INSERT。"""
        with self._connection_factory() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE ck_post SET platform=%s, url=%s, raw=%s, "
                    "updated_at=now() WHERE id=%s",
                    (post.platform, post.url, psycopg2.extras.Json(post.raw), post.id),
                )
                if cur.rowcount == 0:
                    cur.execute(
                        "INSERT INTO ck_post (id, platform, url, raw, stage) "
                        "VALUES (%s, %s, %s, %s, 'fetched')",
                        (post.id, post.platform, post.url,
                         psycopg2.extras.Json(post.raw)),
                    )
            conn.commit()

    def set_extracted(self, post_id: str, extracted: dict, stage: str = "extracted") -> None:
        self._set_json(post_id, "extracted", extracted, stage)

    def set_screening(self, post_id: str, screening: dict, stage: str = "screened") -> None:
        self._set_json(post_id, "screening", screening, stage)

    def _set_json(self, post_id: str, column: str, value: dict, stage: str) -> None:
        with self._connection_factory() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"UPDATE ck_post SET {column}=%s, stage=%s, updated_at=now() "
                    "WHERE id=%s",
                    (psycopg2.extras.Json(value), stage, post_id),
                )
            conn.commit()

    def update_stage(self, post_id: str, stage: str) -> None:
        with self._connection_factory() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE ck_post SET stage=%s, updated_at=now() WHERE id=%s",
                    (stage, post_id),
                )
            conn.commit()

    def save_item(
        self,
        post_id: str,
        item: dict,
        deconstruction: Optional[dict] = None,
        ingest_payload: Optional[dict] = None,
    ) -> int:
        """插入一条知识片段，返回自增 id。ingest_status 默认 pending。"""
        with self._connection_factory() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO ck_knowledge_item "
                    "(post_id, item, deconstruction, ingest_payload, ingest_status) "
                    "VALUES (%s, %s, %s, %s, 'pending') RETURNING id",
                    (
                        post_id,
                        psycopg2.extras.Json(item),
                        psycopg2.extras.Json(deconstruction) if deconstruction else None,
                        psycopg2.extras.Json(ingest_payload) if ingest_payload else None,
                    ),
                )
                item_id = cur.fetchone()[0]
            conn.commit()
            return int(item_id)

    def update_item_ingest(
        self, item_id: int, status: str, knowledge_id: Optional[str] = None
    ) -> None:
        with self._connection_factory() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "UPDATE ck_knowledge_item SET ingest_status=%s, knowledge_id=%s, "
                    "updated_at=now() WHERE id=%s",
                    (status, knowledge_id, item_id),
                )
            conn.commit()

    def delete_post(self, post_id: str) -> None:
        """删除帖子及其片段（用于校验探测的清理）。"""
        with self._connection_factory() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM ck_knowledge_item WHERE post_id=%s", (post_id,))
                cur.execute("DELETE FROM ck_post WHERE id=%s", (post_id,))
            conn.commit()

    # ---------- 读 ----------
    def read_post(self, post_id: str) -> Optional[dict]:
        return self._fetch_one(
            "SELECT * FROM ck_post WHERE id=%s", (post_id,)
        )

    def read_items(self, post_id: str) -> list[dict]:
        return self._fetch_all(
            "SELECT * FROM ck_knowledge_item WHERE post_id=%s ORDER BY id", (post_id,)
        )

    def list_posts(self, limit: int = 100) -> list[dict]:
        return self._fetch_all(
            "SELECT id, platform, url, stage, created_at, updated_at "
            "FROM ck_post ORDER BY created_at DESC LIMIT %s",
            (limit,),
        )

    def table_columns(self, table: str) -> list[str]:
        rows = self._fetch_all(
            "SELECT column_name FROM information_schema.columns "
            "WHERE table_schema=%s AND table_name=%s ORDER BY ordinal_position",
            (self.config.schema, table),
        )
        return [r["column_name"] for r in rows]

    def _fetch_one(self, sql: str, params: tuple) -> Optional[dict]:
        with self._connection_factory() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, params)
                row = cur.fetchone()
                return dict(row) if row else None

    def _fetch_all(self, sql: str, params: tuple) -> list[dict]:
        with self._connection_factory() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(sql, params)
                return [dict(r) for r in cur.fetchall()]