"""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_` 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_` # 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)