Create-knowledge-find-decode/pipeline/postgres.py

"""PostgreSQL repository and best-effort writer for pipeline tracing."""
from __future__ import annotations

import logging
from typing import Any
from uuid import UUID

import psycopg2.extras

from core.config import CreationDbConfig
from core.db_session import transaction
from pipeline.models import PipelineJob, PipelineRun, PipelineRunEvent, PipelineStage
from pipeline.tracing import TraceConfig, TraceContext, sanitize_payload

Json = psycopg2.extras.Json
psycopg2.extras.register_uuid()
logger = logging.getLogger(__name__)


class PostgresPipelineRepository:
    """Repository for orchestration state and trace ledger rows.

    This class owns no transaction boundary; API dependencies pass a pooled
    connection, while the best-effort trace writer opens short transactions.
    """

    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_pipeline_run(
        self,
        *,
        run_key: str | None = None,
        batch_id: UUID | None = None,
        status: str = "pending",
        current_stage: str | None = None,
        config: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> PipelineRun:
        row = self._one(
            """
            INSERT INTO pipeline_runs(
                run_key, batch_id, status, current_stage, config, metadata, started_at
            )
            VALUES (%s, %s, %s, %s, %s, %s, CASE WHEN %s = 'running' THEN now() ELSE NULL END)
            ON CONFLICT (run_key) DO UPDATE SET
                batch_id = COALESCE(EXCLUDED.batch_id, pipeline_runs.batch_id),
                status = EXCLUDED.status,
                current_stage = COALESCE(EXCLUDED.current_stage, pipeline_runs.current_stage),
                config = pipeline_runs.config || EXCLUDED.config,
                metadata = pipeline_runs.metadata || EXCLUDED.metadata,
                started_at = COALESCE(pipeline_runs.started_at, EXCLUDED.started_at)
            RETURNING *
            """,
            (
                run_key,
                batch_id,
                status,
                current_stage,
                Json(config or {}),
                Json(metadata or {}),
                status,
            ),
        )
        return _pipeline_run(row)

    def mark_pipeline_run_status(
        self,
        run_id: UUID,
        *,
        status: str,
        current_stage: str | None = None,
        error_message: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> PipelineRun:
        row = self._one(
            """
            UPDATE pipeline_runs SET
                status = %s,
                current_stage = COALESCE(%s, current_stage),
                error_message = %s,
                metadata = CASE WHEN %s THEN metadata || %s ELSE metadata END,
                finished_at = CASE WHEN %s IN ('done', 'partial', 'failed') THEN now() ELSE finished_at END
            WHERE id = %s
            RETURNING *
            """,
            (
                status,
                current_stage,
                error_message,
                metadata is not None,
                Json(metadata or {}),
                status,
                run_id,
            ),
        )
        return _pipeline_run(row)

    def get_pipeline_run(self, run_id: UUID) -> PipelineRun:
        return _pipeline_run(self._one("SELECT * FROM pipeline_runs WHERE id = %s", (run_id,)))

    def get_pipeline_run_by_acquisition_run(self, acquisition_run_id: UUID) -> PipelineRun:
        row = self._one(
            """
            SELECT pr.* FROM pipeline_runs pr
            JOIN pipeline_run_events pre ON pre.pipeline_run_id = pr.id
            WHERE pre.acquisition_run_id = %s
            ORDER BY pre.created_at DESC
            LIMIT 1
            """,
            (acquisition_run_id,),
        )
        return _pipeline_run(row)

    def save_pipeline_job(
        self,
        *,
        run_id: UUID,
        stage: PipelineStage | str,
        target_id: UUID | None = None,
        target_table: str | None = None,
        status: str = "pending",
        metadata: dict[str, Any] | None = None,
    ) -> PipelineJob:
        row = self._one(
            """
            INSERT INTO pipeline_jobs(
                pipeline_run_id, stage, target_table, target_id, status, metadata, started_at
            )
            VALUES (%s, %s, %s, %s, %s, %s, CASE WHEN %s = 'running' THEN now() ELSE NULL END)
            RETURNING *
            """,
            (run_id, stage, target_table, target_id, status, Json(metadata or {}), status),
        )
        return _pipeline_job(row)

    def mark_job_status(
        self,
        job_id: UUID,
        *,
        status: str,
        error_message: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> PipelineJob:
        row = self._one(
            """
            UPDATE pipeline_jobs SET
                status = %s,
                error_message = %s,
                metadata = CASE WHEN %s THEN metadata || %s ELSE metadata END,
                finished_at = CASE WHEN %s IN ('done', 'partial', 'failed', 'skipped') THEN now() ELSE finished_at END
            WHERE id = %s
            RETURNING *
            """,
            (
                status,
                error_message,
                metadata is not None,
                Json(metadata or {}),
                status,
                job_id,
            ),
        )
        return _pipeline_job(row)

    def append_event(
        self,
        *,
        context: TraceContext,
        stage: str,
        event_type: str,
        status: str | None = None,
        severity: str = "info",
        target_table: str | None = None,
        target_id: UUID | None = None,
        message: str | None = None,
        payload: dict[str, Any] | None = None,
        error_message: str | None = None,
        duration_ms: int | None = None,
        attempt_index: int | None = None,
    ) -> PipelineRunEvent:
        row = self._one(
            """
            INSERT INTO pipeline_run_events(
                pipeline_run_id, pipeline_job_id, stage, event_type, status, severity,
                target_table, target_id, acquisition_run_id, acquisition_job_id,
                query_id, item_id, decode_job_id, decode_result_id, payload_draft_id,
                ingest_record_id, platform, attempt_index, message, payload,
                error_message, duration_ms
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                context.pipeline_run_id,
                context.pipeline_job_id,
                stage,
                event_type,
                status,
                severity,
                target_table,
                target_id,
                context.acquisition_run_id,
                context.acquisition_job_id,
                context.query_id,
                context.item_id,
                context.decode_job_id,
                context.decode_result_id,
                context.payload_draft_id,
                context.ingest_record_id,
                context.platform,
                attempt_index,
                message,
                Json(payload or {}),
                error_message,
                duration_ms,
            ),
        )
        return PipelineRunEvent.model_validate(_event_row(row))

    def record_candidate_hit(
        self,
        *,
        context: TraceContext,
        item_id: UUID | None,
        platform: str,
        unique_key: str | None = None,
        platform_item_id: str | None = None,
        search_provider: str | None = None,
        detail_provider: str | None = None,
        attempt_index: int | None = None,
        page_index: int | None = None,
        page_rank: int | None = None,
        candidate_rank: int | None = None,
        source_cursor: str | None = None,
        is_duplicate_hit: bool = False,
        raw_candidate: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        return self._one(
            """
            INSERT INTO candidate_item_hits(
                pipeline_run_id, acquisition_run_id, acquisition_job_id,
                query_id, item_id, platform, unique_key, platform_item_id,
                search_provider, detail_provider, attempt_index, page_index,
                page_rank, candidate_rank, source_cursor, is_duplicate_hit,
                raw_candidate, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                context.pipeline_run_id,
                context.acquisition_run_id,
                context.acquisition_job_id,
                context.query_id,
                item_id,
                platform,
                unique_key,
                platform_item_id,
                search_provider,
                detail_provider,
                attempt_index,
                page_index,
                page_rank,
                candidate_rank,
                source_cursor,
                is_duplicate_hit,
                Json(raw_candidate or {}),
                Json(metadata or {}),
            ),
        )

    def save_llm_call_trace(
        self,
        *,
        context: TraceContext,
        stage: str,
        substage: str | None = None,
        provider: str | None = None,
        model_name: str | None = None,
        endpoint: str | None = None,
        prompt_name: str | None = None,
        prompt_hash: str | None = None,
        request_payload: dict[str, Any] | None = None,
        response_payload: dict[str, Any] | None = None,
        parsed_payload: dict[str, Any] | None = None,
        status: str = "pending",
        error_message: str | None = None,
        latency_ms: int | None = None,
        attempt_index: int | None = None,
    ) -> dict[str, Any]:
        return self._one(
            """
            INSERT INTO llm_call_traces(
                pipeline_run_id, pipeline_job_id, stage, substage, provider,
                model_name, endpoint, prompt_name, prompt_hash, trace_context,
                request_payload, response_payload, parsed_payload, status,
                error_message, latency_ms, attempt_index
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING *
            """,
            (
                context.pipeline_run_id,
                context.pipeline_job_id,
                stage,
                substage,
                provider,
                model_name,
                endpoint,
                prompt_name,
                prompt_hash,
                Json(context.as_payload()),
                Json(request_payload or {}),
                Json(response_payload or {}),
                Json(parsed_payload or {}),
                status,
                error_message,
                latency_ms,
                attempt_index,
            ),
        )

    def list_timeline(self, run_id: UUID) -> list[dict[str, Any]]:
        return self._all(
            """
            SELECT * FROM pipeline_run_events
            WHERE pipeline_run_id = %s
            ORDER BY created_at, id
            """,
            (run_id,),
        )

    def list_jobs(self, run_id: UUID) -> list[dict[str, Any]]:
        return self._all(
            """
            SELECT * FROM pipeline_jobs
            WHERE pipeline_run_id = %s
            ORDER BY created_at, id
            """,
            (run_id,),
        )

    def list_candidate_hits(self, run_id: UUID) -> list[dict[str, Any]]:
        return self._all(
            """
            SELECT * FROM candidate_item_hits
            WHERE pipeline_run_id = %s
            ORDER BY created_at, page_index NULLS LAST, page_rank NULLS LAST, id
            """,
            (run_id,),
        )

    def list_llm_call_traces(self, run_id: UUID, *, limit: int = 500) -> list[dict[str, Any]]:
        return self._all(
            """
            SELECT * FROM llm_call_traces
            WHERE pipeline_run_id = %s
            ORDER BY created_at, id
            LIMIT %s
            """,
            (run_id, limit),
        )

    def get_timeline_bundle(self, run_id: UUID) -> dict[str, Any]:
        return {
            "events": self.list_timeline(run_id),
            "jobs": self.list_jobs(run_id),
            "candidate_hits": self.list_candidate_hits(run_id),
            "llm_call_traces": self.list_llm_call_traces(run_id),
        }

    def get_resume_cursor(self, run_id: UUID):
        return None


class PostgresTraceWriter:
    """Best-effort writer that never lets tracing break the business flow."""

    def __init__(self, *, db_config: CreationDbConfig, config: TraceConfig):
        self.db_config = db_config
        self.config = config

    def event(self, *, context: TraceContext, stage: str, event_type: str, **kwargs: Any) -> None:
        self._best_effort(
            lambda repo: repo.append_event(
                context=context,
                stage=stage,
                event_type=event_type,
                payload=sanitize_payload(kwargs.pop("payload", {}) or {}, max_chars=self.config.max_payload_chars),
                **kwargs,
            )
        )

    def candidate_hit(self, *, context: TraceContext, **kwargs: Any) -> None:
        self._best_effort(
            lambda repo: repo.record_candidate_hit(
                context=context,
                raw_candidate=sanitize_payload(kwargs.pop("raw_candidate", {}) or {}, max_chars=self.config.max_payload_chars),
                metadata=sanitize_payload(kwargs.pop("metadata", {}) or {}, max_chars=self.config.max_payload_chars),
                **kwargs,
            )
        )

    def llm_call(self, *, context: TraceContext, stage: str, substage: str | None = None, **kwargs: Any) -> None:
        request_payload = kwargs.pop("request_payload", {}) if self.config.capture_request else {}
        response_payload = kwargs.pop("response_payload", {}) if self.config.capture_response else {}
        parsed_payload = kwargs.pop("parsed_payload", {})
        self._best_effort(
            lambda repo: repo.save_llm_call_trace(
                context=context,
                stage=stage,
                substage=substage,
                request_payload=sanitize_payload(request_payload or {}, max_chars=self.config.max_payload_chars),
                response_payload=sanitize_payload(response_payload or {}, max_chars=self.config.max_payload_chars),
                parsed_payload=sanitize_payload(parsed_payload or {}, max_chars=self.config.max_payload_chars),
                **kwargs,
            )
        )

    def _best_effort(self, fn) -> None:
        if not self.config.enabled:
            return
        try:
            with transaction(self.db_config) as conn:
                fn(PostgresPipelineRepository(conn))
        except Exception as exc:  # pragma: no cover - defensive isolation
            logger.warning("pipeline trace write skipped: %s", exc)


def _pipeline_run(row: dict[str, Any]) -> PipelineRun:
    return PipelineRun(
        id=row.get("id"),
        run_key=row.get("run_key"),
        batch_id=row.get("batch_id"),
        status=row.get("status") or "pending",
        current_stage=row.get("current_stage"),
        metadata={
            **(row.get("config") or {}),
            **(row.get("metadata") or {}),
        },
        error_message=row.get("error_message"),
        started_at=row.get("started_at"),
        finished_at=row.get("finished_at"),
    )


def _pipeline_job(row: dict[str, Any]) -> PipelineJob:
    return PipelineJob(
        id=row.get("id"),
        run_id=row.get("pipeline_run_id"),
        stage=row.get("stage"),
        status=row.get("status") or "pending",
        target_table=row.get("target_table"),
        target_id=row.get("target_id"),
        attempt_count=row.get("attempt_count") or 0,
        metadata=row.get("metadata") or {},
        error_message=row.get("error_message"),
        started_at=row.get("started_at"),
        finished_at=row.get("finished_at"),
    )


def _event_row(row: dict[str, Any]) -> dict[str, Any]:
    out = dict(row)
    out["pipeline_run_id"] = out.pop("pipeline_run_id", None)
    return out