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, user_id: str | None = None
    ) -> 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.

        **Tenant scoping (2026-07-23).** When `user_id` is supplied the row must be
        owned by that user. Go enforces exactly this pair on every equivalent read
        (`catalog_repo.go`: `WHERE scope_type='analysis' AND analysis_id=$1 AND
        user_id=$2`); Python filtered on `analysis_id` alone, so a caller who knew
        another tenant's `analysis_id` received that tenant's catalog — and, because
        the payload also carries the owner's `user_id`, `DbExecutor`'s ownership
        check then compared the victim's id against itself and passed, executing SQL
        against their database.

        `user_id` is optional so the legacy/unthreaded call sites keep working, but
        an unscoped read is logged: those call sites are the remaining work.
        """
        async with AsyncSessionLocal() as session:
            where = [
                CatalogRow.analysis_id == analysis_id,
                CatalogRow.scope_type == "analysis",
            ]
            if user_id is not None:
                where.append(CatalogRow.user_id == user_id)
            result = await session.execute(
                select(CatalogRow.catalog_payload).where(*where)
            )
            row = result.scalar_one_or_none()
            if row is None and user_id is not None:
                # Diagnose the miss: an owned row that we just refused is either a
                # genuine cross-tenant attempt or a `user_id` format mismatch between
                # what Go wrote and what the caller sent. Both need to be loud; the
                # answer is the same either way (deny).
                probe = await session.execute(
                    select(CatalogRow.user_id).where(
                        CatalogRow.analysis_id == analysis_id,
                        CatalogRow.scope_type == "analysis",
                    )
                )
                owner = probe.scalar_one_or_none()
                if owner is not None:
                    logger.error(
                        "analysis catalog owner mismatch — denied",
                        analysis_id=analysis_id,
                        requested_by=user_id,
                        owner=owner,
                    )
        if user_id is None:
            logger.warning(
                "analysis catalog read is UNSCOPED (no user_id) — call site needs threading",
                analysis_id=analysis_id,
            )
        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)