src/catalog/reader.py

"""CatalogReader — loads + filters catalog by source_hint.

For typical users (≤50 tables), returns the FULL catalog with no slicing.
Catalog-level search is added later if catalog grows past the limit.
"""

from datetime import UTC, datetime
from typing import Literal

from .models import Catalog, Source
from .store import CatalogStore

SourceHint = Literal["chat", "unstructured", "structured"]


def _filter_sources(catalog: Catalog, source_hint: SourceHint) -> Catalog:
    """Return a copy of `catalog` keeping only the sources matching `source_hint`."""
    filtered: list[Source]
    if source_hint == "chat":
        filtered = []
    elif source_hint == "structured":
        filtered = [s for s in catalog.sources if s.source_type in {"schema", "tabular"}]
    else:  # "unstructured"
        filtered = [s for s in catalog.sources if s.source_type == "unstructured"]
    return catalog.model_copy(update={"sources": filtered})


class CatalogReader:
    """Loads the user's catalog and filters by source_hint.

    On miss, returns an empty Catalog (never raises) — query path is
    responsible for handling "no data registered yet" gracefully.
    Returned Catalog is always a copy; the underlying stored catalog
    is never mutated.
    """

    def __init__(self, store: CatalogStore) -> None:
        self._store = store

    async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
        catalog = await self._store.get(user_id)
        if catalog is None:
            return Catalog(user_id=user_id, generated_at=datetime.now(UTC))

        return _filter_sources(catalog, source_hint)


class MemoizingCatalogReader(CatalogReader):
    """Request-scoped CatalogReader that caches each ``read`` by source_hint.

    One per request. The same per-user catalog is otherwise fetched from the
    catalog DB 4-5x during a single slow-path run (planner load, then
    check_data's structured read + check_knowledge's unstructured read, then
    retrieve_data's structured read). Wrapping the base reader collapses those
    to one round-trip
    per distinct source_hint and pins a single consistent snapshot for the whole
    request (plan-time and execution-time catalogs can no longer diverge).
    """

    def __init__(self, inner: CatalogReader) -> None:
        # `read` is fully overridden below and delegates to `inner`, so the parent's
        # `_store` is never used — carry it through only so this stays a real
        # CatalogReader (any inner with a `read` works, including test fakes).
        super().__init__(getattr(inner, "_store", None))
        self._inner = inner
        self._cache: dict[SourceHint, Catalog] = {}

    async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
        cached = self._cache.get(source_hint)
        if cached is None:
            cached = await self._inner.read(user_id, source_hint)
            self._cache[source_hint] = cached
        return cached


class AnalysisScopedCatalogReader(CatalogReader):
    """Reads the analysis-scope catalog, falling back to the user-scope reader.

    Used by the `check` skill so "what data do I have" inside a room reflects
    that analysis's bound sources — structured AND documents — with their real
    names. A database shows as "xl test" (analysis-scope) instead of the
    auto-generated `postgres_<hash>` placeholder, and documents show at all
    (the user-scope catalog holds no `unstructured` sources, so reading them from
    user-scope always came back empty). When the analysis has no catalog row
    (legacy / not yet bound) or the read fails, it degrades to the wrapped
    user-scope reader, so unbound rooms behave exactly as before.
    """

    def __init__(self, inner: CatalogReader, analysis_id: str | None) -> None:
        # `inner` is a real CatalogReader (constructed at the check call site), so
        # its `_store` is the live CatalogStore we need for the analysis read.
        super().__init__(inner._store)
        self._inner = inner
        self._analysis_id = analysis_id

    async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
        # Read analysis-scope for BOTH structured and unstructured. Verified via
        # the dedorch `data_catalog` table: the analysis-scope rows carry the real
        # DB names AND the room's documents (`source_type='unstructured'`), whereas
        # the user-scope rows hold only structured sources with `postgres_<hash>`
        # placeholder names and NO documents at all — so reading documents from
        # user-scope always returned empty ("not listed"). Fall back to the
        # user-scope reader only when the analysis has no catalog row (legacy /
        # unbound room).
        if self._analysis_id:
            try:
                catalog = await self._store.get_by_analysis(self._analysis_id)
            except Exception:  # noqa: BLE001 — never block check on the analysis read
                catalog = None
            if catalog is not None:
                return _filter_sources(catalog, source_hint)
        return await self._inner.read(user_id, source_hint)