src/agents/slow_path/store.py

"""ReportInputStore — the seam the slow path persists its AnalysisRecord through.

The Assembler produces an `AnalysisRecord` (the faithful, structured record of a
run — §8.3, INV-4). Persisting it is a separate concern from streaming the answer,
so it sits behind this seam. `generate_report` later reads records back by
`analysis_id` (oldest-first) and renders from them — never from chat history.

- `NullReportInputStore` logs and stores nothing (kept for tests / when persistence
  is intentionally disabled).
- `PostgresReportInputStore` writes one `report_inputs` row per run in the catalog
  DB (Neon `dataeyond`, `settings.postgres_connstring`).

`save` must never raise on the caller's path — a persistence failure must not break
the user's answer (§8.3). `list_for_analysis` is a read for the report generator and
is allowed to surface errors to its caller.
"""

from __future__ import annotations

from typing import Protocol, runtime_checkable

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert

from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import ReportInputRow
from src.middlewares.logging import get_logger

from .schemas import AnalysisRecord

logger = get_logger("analysis_store")


@runtime_checkable
class ReportInputStore(Protocol):
    """Persist + read completed analyses.

    `save` must never raise on the caller's path. `list_for_analysis` returns the
    records for one analysis session, oldest-first (the order the report renders in).
    """

    async def save(self, record: AnalysisRecord) -> None: ...

    async def list_for_analysis(self, analysis_id: str) -> list[AnalysisRecord]: ...


class NullReportInputStore:
    """No-op store: logs the record, persists nothing. Reads return empty."""

    async def save(self, record: AnalysisRecord) -> None:
        logger.info(
            "analysis_record produced (not persisted — NullReportInputStore)",
            record_id=record.record_id,
            plan_id=record.plan_id,
            n_tasks=len(record.tasks_run),
        )

    async def list_for_analysis(self, analysis_id: str) -> list[AnalysisRecord]:
        return []


class PostgresReportInputStore:
    """Writes/reads `report_inputs` jsonb rows in the catalog DB.

    Mirrors `CatalogStore`: each call opens its own `AsyncSession`. One row per
    record (vs. one-per-user for the catalog) since records accumulate per analysis.
    """

    async def save(self, record: AnalysisRecord) -> None:
        try:
            payload = record.model_dump(mode="json")
            async with AsyncSessionLocal() as session:
                stmt = insert(ReportInputRow).values(
                    id=record.record_id,
                    analysis_id=record.analysis_id,
                    user_id=record.user_id,
                    plan_id=record.plan_id,
                    data=payload,
                    created_at=record.created_at,
                )
                # Re-running the same plan id-collides only if record_id repeats;
                # treat that as idempotent (overwrite) rather than erroring the user.
                stmt = stmt.on_conflict_do_update(
                    index_elements=[ReportInputRow.id],
                    set_={"data": stmt.excluded.data},
                )
                await session.execute(stmt)
                await session.commit()
            logger.info(
                "analysis_record persisted",
                record_id=record.record_id,
                analysis_id=record.analysis_id,
                user_id=record.user_id,
            )
        except Exception as exc:  # never break the user's answer (§8.3)
            logger.error(
                "analysis_record persist failed",
                record_id=record.record_id,
                error=str(exc),
            )

    async def list_for_analysis(self, analysis_id: str) -> list[AnalysisRecord]:
        async with AsyncSessionLocal() as session:
            result = await session.execute(
                select(ReportInputRow.data)
                .where(ReportInputRow.analysis_id == analysis_id)
                .order_by(ReportInputRow.created_at.asc())
            )
            rows = result.scalars().all()
        return [AnalysisRecord.model_validate(row) for row in rows]