src/charts/store.py

"""ChartStore — the seam the chat pipeline persists `render_chart` outputs through
(S2 visualization, SPINE_V2_PLAN §4.4).

`ChatHandler._run_slow_path` scans a completed `AnalysisRecord.results_snapshot` for
`ToolOutput(kind="chart")` entries and flushes one `MessageChartRow` per chart through
this seam, right before the `done` SSE event (mirrors the traceability flush). Unlike
traceability's one row per turn, a turn with multiple `render_chart` calls writes
multiple rows sharing one `message_id`. `GET /api/v1/charts` reads them back by
(analysis_id, message_id).

- `NullChartStore` logs the envelope and stores nothing (tests / disabled persistence).
- `PostgresChartStore` writes one `message_charts` row per chart (dedorch,
  `AsyncSessionLocal`), mirroring `PostgresTraceabilityStore`.

`save` must NEVER raise on the caller's path — a chart-persist failure must not break
the user's answer. `list_for_message` is the endpoint read; an empty list is a valid
result (chartless turn), not an error.
"""

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Protocol, runtime_checkable

from pydantic import BaseModel
from sqlalchemy import select

from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import MessageChartRow, MessageTraceabilityRow
from src.middlewares.logging import get_logger

logger = get_logger("charts_store")


class ChartRecord(BaseModel):
    """One persisted chart, as served by `GET /api/v1/charts`."""

    chart_id: str
    chart_type: str
    title: str | None = None
    spec: dict
    created_at: datetime


@runtime_checkable
class ChartStore(Protocol):
    """Persist + read `render_chart` outputs for one assistant `message_id`.

    `save` must never raise on the caller's path. `list_for_message` returns every
    chart for one turn (possibly empty — a turn need not have asked for a chart);
    the lookup is by `message_id` alone (Python-minted UUID4, globally unique —
    lead decision 2026-07-13). `turn_exists` says whether the turn flushed a
    traceability row, so the endpoint can tell "no charts" from "unknown id".
    """

    async def save(
        self,
        *,
        message_id: str,
        analysis_id: str,
        user_id: str,
        record_id: str | None,
        envelope: dict,
    ) -> None: ...

    async def list_for_message(self, message_id: str) -> list[ChartRecord]: ...

    async def turn_exists(self, message_id: str) -> bool: ...


class NullChartStore:
    """No-op store: logs the envelope, persists nothing. Reads return `[]`."""

    async def save(
        self,
        *,
        message_id: str,
        analysis_id: str,
        user_id: str,
        record_id: str | None,
        envelope: dict,
    ) -> None:
        logger.info(
            "chart produced (not persisted — NullChartStore)",
            message_id=message_id,
            chart_type=envelope.get("chart_type", "unknown"),
        )

    async def list_for_message(self, message_id: str) -> list[ChartRecord]:
        return []

    async def turn_exists(self, message_id: str) -> bool:
        return False


class PostgresChartStore:
    """Writes/reads `message_charts` rows. One insert per chart (not an upsert)."""

    async def save(
        self,
        *,
        message_id: str,
        analysis_id: str,
        user_id: str,
        record_id: str | None,
        envelope: dict,
    ) -> None:
        try:
            async with AsyncSessionLocal() as session:
                row = MessageChartRow(
                    id=str(uuid.uuid4()),
                    message_id=message_id,
                    analysis_id=analysis_id,
                    user_id=user_id,
                    record_id=record_id,
                    chart_type=envelope.get("chart_type", "unknown"),
                    title=envelope.get("title"),
                    spec=envelope,
                )
                session.add(row)
                await session.commit()
            logger.info(
                "chart persisted",
                message_id=message_id,
                analysis_id=analysis_id,
                chart_type=envelope.get("chart_type", "unknown"),
            )
        except Exception as exc:  # never break the user's answer
            logger.error(
                "chart persist failed",
                degraded_seam="chart_persist",
                message_id=message_id,
                error=repr(exc),
            )

    async def list_for_message(self, message_id: str) -> list[ChartRecord]:
        # message_id-only lookup (lead decision 2026-07-13). NOTE for the Harry
        # migration: the manual DDL's composite index (analysis_id, message_id)
        # does not serve this predicate — an additive index on (message_id) is
        # part of the handoff.
        async with AsyncSessionLocal() as session:
            result = await session.execute(
                select(MessageChartRow)
                .where(MessageChartRow.message_id == message_id)
                .order_by(MessageChartRow.created_at)
            )
            rows = result.scalars().all()
        return [
            ChartRecord(
                chart_id=row.id,
                chart_type=row.chart_type,
                title=row.title,
                spec=row.spec,
                created_at=row.created_at,
            )
            for row in rows
        ]

    async def turn_exists(self, message_id: str) -> bool:
        """True iff the turn flushed its traceability row (written before `done`).

        Lets the endpoint tri-state a zero-chart GET: `empty` (completed turn, no
        charts — the common case) vs `not_found` (unknown/mistyped id, or an error
        turn, which never writes traceability). PK lookup — cheap.
        """
        async with AsyncSessionLocal() as session:
            result = await session.execute(
                select(MessageTraceabilityRow.message_id).where(
                    MessageTraceabilityRow.message_id == message_id
                )
            )
            return result.scalar_one_or_none() is not None