File size: 4,348 Bytes
0e02a0f
81e5fe7
 
 
0721bb4
 
81e5fe7
0e02a0f
0721bb4
0e02a0f
0721bb4
81e5fe7
0721bb4
 
 
81e5fe7
 
 
 
 
 
0721bb4
 
 
 
0e02a0f
81e5fe7
 
 
 
 
 
 
 
0e02a0f
0721bb4
 
 
 
 
81e5fe7
 
 
0721bb4
 
81e5fe7
0e02a0f
0721bb4
81e5fe7
 
 
0e02a0f
0721bb4
81e5fe7
 
 
0721bb4
 
 
 
 
0e02a0f
 
0721bb4
 
 
 
 
 
 
 
 
0e02a0f
0721bb4
 
 
 
 
 
 
 
 
 
0e02a0f
0721bb4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e02a0f
 
 
0721bb4
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
"""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]