Create-knowledge-find-decode/decode_content/repositories/postgres.py

"""PostgreSQL implementation for formal decode-content state."""
from __future__ import annotations

from typing import Any
from uuid import UUID

import psycopg2.extras

from decode_content.models import (
    ContractSnapshot,
    DecodeJob,
    DecodeResult,
    IngestRecord,
    KnowledgeParticle,
    PayloadDraft,
    ScopeResult,
)

Json = psycopg2.extras.Json
psycopg2.extras.register_uuid()


def _json(value: dict[str, Any] | None) -> Json:
    return Json(value or {})


def _stage_value(value: Any) -> str | None:
    if isinstance(value, list):
        return " / ".join(str(item) for item in value if item)
    if value:
        return str(value)
    return None


class PostgresDecodeRepository:
    """Repository backed by the formal PostgreSQL schema.

    Like the acquisition repository, this class owns no transaction boundary;
    callers pass an open connection and commit/rollback outside.
    """

    def __init__(self, conn: Any):
        self.conn = conn

    def _one(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any]:
        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            row = cur.fetchone()
            if row is None:
                raise RuntimeError("expected one row, got none")
            return dict(row)

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

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

    def create_decode_job(
        self,
        *,
        item_id: UUID,
        status: str = "pending",
        metadata: dict[str, Any] | None = None,
    ) -> DecodeJob:
        row = self._one(
            """
            INSERT INTO decode_jobs(item_id, status, metadata, started_at)
            VALUES (%s, %s, %s, CASE WHEN %s = 'running' THEN now() ELSE NULL END)
            RETURNING *
            """,
            (item_id, status, _json(metadata), status),
        )
        return DecodeJob.model_validate(row)

    def save_decode_result(
        self,
        *,
        item_id: UUID,
        decode_job_id: UUID | None = None,
        read_result: dict[str, Any] | None = None,
        gate_result: dict[str, Any] | None = None,
        framing_result: dict[str, Any] | None = None,
        status: str = "draft",
    ) -> DecodeResult:
        row = self._one(
            """
            INSERT INTO decode_results(
                decode_job_id, item_id, read_result, gate_result,
                framing_result, status
            )
            VALUES (%s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                decode_job_id,
                item_id,
                _json(read_result),
                _json(gate_result),
                _json(framing_result),
                status,
            ),
        )
        if decode_job_id:
            self._one(
                """
                UPDATE decode_jobs
                SET status = %s, finished_at = now(), error_message = NULL
                WHERE id = %s
                RETURNING *
                """,
                (status, decode_job_id),
            )
        return DecodeResult.model_validate(row)

    def get_decode_result_for_item(self, item_id: UUID) -> DecodeResult:
        row = self._one(
            """
            SELECT * FROM decode_results
            WHERE item_id = %s
            ORDER BY created_at DESC
            LIMIT 1
            """,
            (item_id,),
        )
        return DecodeResult.model_validate(row)

    def mark_running_decode_jobs_failed(self, item_id: UUID, error_message: str) -> None:
        with self.conn.cursor() as cur:
            cur.execute(
                """
                UPDATE decode_jobs
                SET status = 'failed', finished_at = now(), error_message = %s
                WHERE item_id = %s AND status = 'running'
                """,
                (error_message, item_id),
            )

    def save_knowledge_particle(
        self,
        *,
        item_id: UUID,
        particle_type: str,
        title: str,
        decode_result_id: UUID | None = None,
        parent_particle_id: UUID | None = None,
        content: dict[str, Any] | None = None,
        status: str = "draft",
    ) -> KnowledgeParticle:
        content = content or {}
        creation_stage = _stage_value(content.get("创作阶段"))
        if creation_stage is None:
            for step in content.get("steps") or []:
                creation_stage = _stage_value(step.get("创作阶段"))
                if creation_stage:
                    break
        row = self._one(
            """
            INSERT INTO knowledge_particles(
                decode_result_id, item_id, parent_particle_id, particle_type,
                title, business_stage, creation_stage, content, status
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                decode_result_id,
                item_id,
                parent_particle_id,
                particle_type,
                title,
                _stage_value(content.get("业务阶段")),
                creation_stage,
                _json(content),
                status,
            ),
        )
        return KnowledgeParticle.model_validate(row)

    def list_knowledge_particles(self, item_id: UUID | None = None) -> list[KnowledgeParticle]:
        if item_id is None:
            rows = self._all(
                "SELECT * FROM knowledge_particles ORDER BY created_at DESC",
                (),
            )
        else:
            rows = self._all(
                """
                SELECT * FROM knowledge_particles
                WHERE item_id = %s
                ORDER BY sort_order, created_at
                """,
                (item_id,),
            )
        return [KnowledgeParticle.model_validate(row) for row in rows]

    def save_scope_result(
        self,
        *,
        scope_type: str,
        scope_value: str,
        particle_id: UUID | None = None,
        item_id: UUID | None = None,
        evidence: dict[str, Any] | None = None,
        status: str = "draft",
    ) -> ScopeResult:
        evidence = evidence or {}
        row = self._one(
            """
            INSERT INTO scope_results(
                particle_id, item_id, scope_type, scope_value, is_reused,
                matched_scope_id, confidence, evidence, status
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                particle_id,
                item_id,
                scope_type,
                scope_value,
                evidence.get("is_reused"),
                evidence.get("matched_scope_id"),
                evidence.get("confidence"),
                _json(evidence),
                status,
            ),
        )
        return ScopeResult.model_validate(row)

    def list_scope_results(self, item_id: UUID | None = None) -> list[ScopeResult]:
        if item_id is None:
            rows = self._all("SELECT * FROM scope_results ORDER BY created_at DESC", ())
        else:
            rows = self._all(
                """
                SELECT * FROM scope_results
                WHERE item_id = %s
                ORDER BY created_at
                """,
                (item_id,),
            )
        return [ScopeResult.model_validate(row) for row in rows]

    def save_payload_draft(
        self,
        *,
        payload: dict[str, Any],
        particle_id: UUID | None = None,
        item_id: UUID | None = None,
        review_status: str = "pending",
        ingest_ready: bool = False,
        status: str = "draft",
    ) -> PayloadDraft:
        row = self._one(
            """
            INSERT INTO payload_drafts(
                particle_id, item_id, payload, review_status, ingest_ready, status
            )
            VALUES (%s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (particle_id, item_id, _json(payload), review_status, ingest_ready, status),
        )
        return PayloadDraft.model_validate(row)

    def list_payload_drafts(self, item_id: UUID | None = None) -> list[PayloadDraft]:
        if item_id is None:
            rows = self._all(
                "SELECT * FROM payload_drafts ORDER BY created_at DESC",
                (),
            )
        else:
            rows = self._all(
                """
                SELECT * FROM payload_drafts
                WHERE item_id = %s
                ORDER BY created_at
                """,
                (item_id,),
            )
        return [PayloadDraft.model_validate(row) for row in rows]

    def mark_payload_draft_ingested(self, payload_draft_id: UUID) -> PayloadDraft:
        row = self._one(
            """
            UPDATE payload_drafts
            SET review_status = 'approved',
                ingest_ready = true,
                status = 'ingested',
                error_message = NULL
            WHERE id = %s
            RETURNING *
            """,
            (payload_draft_id,),
        )
        return PayloadDraft.model_validate(row)

    def save_ingest_record(
        self,
        *,
        payload_draft_id: UUID | None = None,
        target_system: str,
        target_id: str | None = None,
        status: str = "pending",
        attempt_count: int = 1,
        response_payload: dict[str, Any] | None = None,
        error_message: str | None = None,
    ) -> IngestRecord:
        row = self._one(
            """
            INSERT INTO ingest_records(
                payload_draft_id, target_system, target_id, status,
                attempt_count, response_payload, error_message
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                payload_draft_id,
                target_system,
                target_id,
                status,
                attempt_count,
                _json(response_payload),
                error_message,
            ),
        )
        return IngestRecord.model_validate(row)

    def list_ingest_records(self, item_id: UUID | None = None) -> list[IngestRecord]:
        if item_id is None:
            rows = self._all("SELECT * FROM ingest_records ORDER BY created_at DESC", ())
        else:
            rows = self._all(
                """
                SELECT ir.* FROM ingest_records ir
                JOIN payload_drafts pd ON pd.id = ir.payload_draft_id
                WHERE pd.item_id = %s
                ORDER BY ir.created_at
                """,
                (item_id,),
            )
        return [IngestRecord.model_validate(row) for row in rows]

    def save_contract_snapshot(
        self,
        *,
        contract_name: str,
        contract_type: str,
        version_label: str | None = None,
        content_hash: str | None = None,
        source_path: str | None = None,
        snapshot: dict[str, Any] | None = None,
    ) -> ContractSnapshot:
        row = self._one(
            """
            INSERT INTO contract_snapshots(
                contract_name, contract_type, version_label, content_hash,
                source_path, snapshot
            )
            VALUES (%s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                contract_name,
                contract_type,
                version_label,
                content_hash,
                source_path,
                _json(snapshot),
            ),
        )
        return ContractSnapshot.model_validate(row)