src/catalog/store.py

"""CatalogStore — reads the per-user catalog from the dedorch `data_catalog` table.

Storage shape (Go-owned): one row per scope in `data_catalog`
(id, scope_type, user_id, analysis_id, catalog_payload jsonb, schema_version,
generated_at, updated_at). Python reads the user-scoped row (scope_type='user');
Go's `catalog.Service` owns all writes, so `upsert`/`remove_source` are legacy.
"""

from sqlalchemy import case, delete, func, select
from sqlalchemy.dialects.postgresql import insert

from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import Catalog as CatalogRow
from src.middlewares.logging import get_logger

from .fk_inference import infer_foreign_keys
from .models import Catalog
from .sample_decode import decode_sample_values

logger = get_logger("catalog_store")


class CatalogStore:
    """Read/write catalogs keyed by user_id.

    Each method opens its own AsyncSession. Callers needing transactional
    coordination across multiple stores can be refactored to accept an
    explicit AsyncSession in a later PR.
    """

    async def get(self, user_id: str) -> Catalog | None:
        async with AsyncSessionLocal() as session:
            result = await session.execute(
                select(CatalogRow.catalog_payload).where(
                    CatalogRow.user_id == user_id,
                    CatalogRow.scope_type == "user",
                )
            )
            row = result.scalar_one_or_none()
        if row is None:
            return None
        # dedorch catalogs ship no foreign_keys (Go introspection drops them),
        # but the IR validator only allows FK-backed joins. Infer the obvious
        # edges so the planner and validator agree. No-op once Go emits real FKs.
        catalog = infer_foreign_keys(Catalog.model_validate(row))
        # dedorch also JSON-marshals numeric sample bytes as base64 (Go bug) —
        # decode them so the planner sees value ranges, not gibberish.
        # No-op once Go emits plain numeric samples.
        decode_sample_values(catalog)
        return catalog

    async def get_by_analysis(self, analysis_id: str) -> Catalog | None:
        """Read the `scope_type='analysis'` catalog row for an analysis.

        Distinct from `get()` (which reads the user-scope row): the analysis-scope
        payload carries the sources actually bound to this analysis AND their
        real names (a database is named e.g. "xl test" here, vs the auto-generated
        `postgres_<hash>` placeholder in the user-scope row). Returns None when the
        analysis has no catalog row (legacy / not yet bound), so callers fall back
        to the user-scope catalog.
        """
        async with AsyncSessionLocal() as session:
            result = await session.execute(
                select(CatalogRow.catalog_payload).where(
                    CatalogRow.analysis_id == analysis_id,
                    CatalogRow.scope_type == "analysis",
                )
            )
            row = result.scalar_one_or_none()
        if row is None:
            return None
        catalog = infer_foreign_keys(Catalog.model_validate(row))
        decode_sample_values(catalog)
        return catalog

    async def upsert(self, catalog: Catalog) -> None:
        # Legacy: Go's catalog.Service owns catalog writes now. Kept working (and
        # reconciled to the dedorch shape) but no longer on any live Python path.
        payload = catalog.model_dump(mode="json")
        async with AsyncSessionLocal() as session:
            stmt = insert(CatalogRow).values(
                scope_type="user",
                user_id=catalog.user_id,
                catalog_payload=payload,
                schema_version=catalog.schema_version,
                generated_at=catalog.generated_at,
                updated_at=func.now(),
            )
            stmt = stmt.on_conflict_do_update(
                index_elements=[CatalogRow.user_id],
                index_where=CatalogRow.scope_type == "user",
                set_={
                    "catalog_payload": stmt.excluded.catalog_payload,
                    "schema_version": stmt.excluded.schema_version,
                    "updated_at": case(
                        (
                            stmt.excluded.catalog_payload != CatalogRow.catalog_payload,
                            func.now(),
                        ),
                        else_=CatalogRow.updated_at,
                    ),
                },
            )
            await session.execute(stmt)
            await session.commit()
        logger.info(
            "catalog upserted",
            user_id=catalog.user_id,
            sources=len(catalog.sources),
        )

    async def remove_source(self, user_id: str, source_id: str) -> None:
        existing = await self.get(user_id)
        if existing is None:
            logger.info("remove_source: no catalog found", user_id=user_id, source_id=source_id)
            return
        filtered = [s for s in existing.sources if s.source_id != source_id]
        if len(filtered) == len(existing.sources):
            logger.info(
                "remove_source: source not in catalog", user_id=user_id, source_id=source_id
            )
            return
        await self.upsert(existing.model_copy(update={"sources": filtered}))
        logger.info("remove_source: source removed", user_id=user_id, source_id=source_id)

    async def delete(self, user_id: str) -> None:
        async with AsyncSessionLocal() as session:
            await session.execute(delete(CatalogRow).where(CatalogRow.user_id == user_id))
            await session.commit()
        logger.info("catalog deleted", user_id=user_id)