Published tables¶

class CatalogAPI:
    """Catalog-layer read API — brokered batches from published catalog assets.

    Mounted at ``AsyncNexus.catalog``. Implements session/run/delivery;
    catalog semantic layer per ; consumes resolved refs per
    (no silent collection-only resolution).
    """

    def __init__(self, nexus_client: NexusClient, *, progress_settings: ProgressSettings | None = None) -> None:
        self._nexus = nexus_client
        self._progress_settings = resolve_progress_settings(progress_settings)

    def asset(
        self,
        ref: CatalogAssetRef | CatalogAssetSnapshotRef | CatalogAsset,
    ) -> CatalogAsset:
        """Hydrate a catalog ref (or pass through an existing :class:`CatalogAsset`)."""
        from miura.nexus.catalog.catalog_asset import CatalogAsset

        if isinstance(ref, CatalogAsset):
            return ref
        return CatalogAsset.from_ref(as_catalog_asset_ref(ref), catalog=self)

    def _control_client(self, *, project_id: str) -> CatalogReadControlClient:
        identity = resolve_catalog_read_identity(project_id=project_id)
        return CatalogReadControlClient(
            catalog_read_base_url(self._nexus.base_url),
            self._nexus.http_client,
            identity=identity,
        )

    async def create_read_session(
        self,
        *,
        asset_id: str = "",
        collection_id: str = "",
        project_id: str = "",
        columns: list[str] | None = None,
        filters: list[dict[str, Any]] | None = None,
        asset_ref: CatalogAssetRef | CatalogAssetSnapshotRef | CatalogAsset | None = None,
    ) -> CatalogReadSession:
        """Advanced control-plane session."""
        ref = resolve_catalog_asset_ref(
            asset_ref=asset_ref,
            asset_id=asset_id,
            collection_id=collection_id,
            project_id=project_id,
        )
        control = self._control_client(project_id=ref.project_id)
        session = await control.create_session(
            collection_id=ref.collection_id,
            asset_id=ref.asset_id,
            requested_columns=columns,
            requested_filters=filters,
        )
        return CatalogReadSession(
            catalog=self,
            session_id=session["session_id"],
            asset_id=ref.asset_id,
            collection_id=ref.collection_id,
            project_id=ref.project_id,
        )

    async def iter_batches(
        self,
        *,
        asset_id: str = "",
        collection_id: str = "",
        project_id: str = "",
        columns: list[str] | None = None,
        filters: list[dict[str, Any]] | None = None,
        batch_size: int = 4096,
        shuffle: dict[str, Any] | None = None,
        shard: dict[str, Any] | CatalogReadShard | None = None,
        asset_ref: CatalogAssetRef | CatalogAssetSnapshotRef | CatalogAsset | None = None,
    ) -> AsyncIterator[pa.RecordBatch]:
        """Neutral catalog batch iterator — control API + Flight delivery."""
        ref = resolve_catalog_asset_ref(
            asset_ref=asset_ref,
            asset_id=asset_id,
            collection_id=collection_id,
            project_id=project_id,
        )
        control = self._control_client(project_id=ref.project_id)
        shard_policy = shard.as_policy() if isinstance(shard, CatalogReadShard) else shard
        from miura.nexus.catalog.catalog_read_progress import CatalogReadProgressTracker

        progress = CatalogReadProgressTracker(progress_settings=self._progress_settings)
        try:
            async for batch in iter_batches_via_transport(
                control,
                collection_id=ref.collection_id,
                asset_id=ref.asset_id,
                columns=columns,
                filters=filters,
                batch_size=batch_size,
                shuffle_policy=shuffle,
                shard_policy=shard_policy,
            ):
                progress.record_batch(rows=batch.num_rows)
                yield batch
        finally:
            progress.close()

    async def iter_batches_for_run(
        self,
        *,
        project_id: str,
        session_id: str,
        run_id: str,
    ) -> AsyncIterator[pa.RecordBatch]:
        """Consume batches for an existing run (advanced session → run path)."""
        control = self._control_client(project_id=project_id)
        from miura.nexus.catalog.catalog_read_progress import CatalogReadProgressTracker

        progress = CatalogReadProgressTracker(progress_settings=self._progress_settings)
        try:
            async for batch in iter_batches_from_delivery(
                control,
                session_id=session_id,
                run_id=run_id,
            ):
                progress.record_batch(rows=batch.num_rows)
                yield batch
        finally:
            progress.close()

class CatalogAsset(CatalogScopedEntity):
    """Queryable published catalog table — catalog-layer entity with read methods.

    Identity and serialization: :attr:`ref` (``CatalogAssetSnapshotRef`` or ``CatalogAssetRef``).
    Behavior: :meth:`iter_batches`, adapter helpers, optional advanced session API.
    """

    __slots__ = ("_ref",)

    def __init__(
        self,
        *,
        catalog: CatalogAPI,
        ref: CatalogAssetRef | CatalogAssetSnapshotRef,
    ) -> None:
        super().__init__(
            catalog=catalog,
            asset_id=ref.asset_id,
            collection_id=ref.collection_id,
            project_id=ref.project_id,
        )
        self._ref = ref

    @property
    def ref(self) -> CatalogAssetRef | CatalogAssetSnapshotRef:
        """Immutable identity handle for serialization, jobs, and low-level APIs."""
        return self._ref

    @property
    def logical_ref(self) -> CatalogAssetRef:
        if isinstance(self._ref, CatalogAssetSnapshotRef):
            return self._ref.logical_ref
        return self._ref

    @property
    def fileset_snapshot_id(self) -> str | None:
        if isinstance(self._ref, CatalogAssetSnapshotRef):
            return self._ref.fileset_snapshot_id
        return None

    @property
    def catalog_registration_id(self) -> str | None:
        if isinstance(self._ref, CatalogAssetSnapshotRef):
            return self._ref.catalog_registration_id
        return None

    @property
    def snapshot_id(self) -> str | None:
        """Alias for :attr:`fileset_snapshot_id` when snapshot-bound."""
        return self.fileset_snapshot_id

    @classmethod
    def from_ref(
        cls,
        ref: CatalogAssetRef | CatalogAssetSnapshotRef,
        *,
        catalog: CatalogAPI,
    ) -> CatalogAsset:
        """Hydrate a ref into a catalog entity (mid-level entry)."""
        return cls(catalog=catalog, ref=ref)

    @classmethod
    def from_nexus_client(
        cls,
        ref: CatalogAssetRef | CatalogAssetSnapshotRef,
        nexus_client: NexusClient,
    ) -> CatalogAsset:
        from miura.nexus.catalog.catalog_api import CatalogAPI

        return cls.from_ref(ref, catalog=CatalogAPI(nexus_client))

    @classmethod
    def from_transport(
        cls,
        ref: CatalogAssetRef | CatalogAssetSnapshotRef,
        transport: Any,
    ) -> CatalogAsset:
        """Wrap a ref using an object exposing ``http_client`` and ``base_url``."""
        return cls.from_ref(ref, catalog=catalog_api_for_transport(transport))

    async def iter_batches(
        self,
        *,
        columns: list[str] | None = None,
        filters: list[dict[str, Any]] | None = None,
        batch_size: int = 4096,
        shuffle: dict[str, Any] | None = None,
        shard: dict[str, Any] | CatalogReadShard | None = None,
    ) -> AsyncIterator[pa.RecordBatch]:
        """Stream record batches for this asset (delegates to :meth:`CatalogAPI.iter_batches`)."""
        async for batch in self._catalog.iter_batches(
            asset_ref=self._ref,
            columns=columns,
            filters=filters,
            batch_size=batch_size,
            shuffle=shuffle,
            shard=shard,
        ):
            yield batch

    async def create_read_session(
        self,
        *,
        columns: list[str] | None = None,
        filters: list[dict[str, Any]] | None = None,
    ) -> CatalogReadSession:
        """Advanced control-plane session for this asset."""
        return await self._catalog.create_read_session(
            asset_ref=self._ref,
            columns=columns,
            filters=filters,
        )

    def to_dataset_source(
        self,
        *,
        columns: Sequence[str] | None = None,
        catalog_batch_size: int = 4096,
        shard_world_size: int = 1,
    ) -> DatasetSource:
        from miura.adapters.dataset_source import DatasetSource

        return DatasetSource.from_catalog_asset(
            self,
            columns=columns,
            catalog_batch_size=catalog_batch_size,
            shard_world_size=shard_world_size,
        )

    def to_torch_dataloader(self, **kwargs: Any) -> Any:
        from miura.adapters.torch import to_torch_dataloader

        return to_torch_dataloader(self, **kwargs)

    def batch_stream_params(
        self,
        *,
        columns: Sequence[str] | None = None,
        catalog_batch_size: int = 4096,
        shuffle: Mapping[str, Any] | None = None,
        shard: CatalogReadShard | None = None,
    ):
        from miura.adapters.catalog_batch_stream import CatalogBatchStreamParams

        return CatalogBatchStreamParams.from_asset_ref(
            self,
            columns=columns,
            catalog_batch_size=catalog_batch_size,
            shuffle=dict(shuffle) if shuffle is not None else None,
            shard=shard,
        )

class CatalogReadRun:
    """Advanced run handle — batch iteration without exposing Flight mechanics."""

    def __init__(
        self,
        *,
        session: CatalogReadSession,
        run_id: str,
        delivery: dict[str, Any],
    ) -> None:
        self._session = session
        self.run_id = run_id
        self._delivery = delivery

    @property
    def session_id(self) -> str:
        return self._session.session_id

    async def iter_batches(self) -> AsyncIterator[pa.RecordBatch]:
        async for batch in self._session._catalog.iter_batches_for_run(
            project_id=self._session.project_id,
            session_id=self.session_id,
            run_id=self.run_id,
        ):
            yield batch

class CatalogScopedEntity:
    """Catalog-layer object bound to a project/collection/asset scope and :class:`CatalogAPI`.

    Subclasses (``CatalogAsset``, ``CatalogReadSession``, …) share identity accessors
    and a single catalog client — behavior differs by entity role, not by duplicating scope.
    """

    __slots__ = ("_catalog", "_asset_id", "_collection_id", "_project_id")

    def __init__(
        self,
        *,
        catalog: CatalogAPI,
        asset_id: str,
        collection_id: str,
        project_id: str,
    ) -> None:
        self._catalog = catalog
        self._asset_id = asset_id
        self._collection_id = collection_id
        self._project_id = project_id

    @property
    def catalog(self) -> CatalogAPI:
        """Low-level catalog namespace (``iter_batches``, sessions, …)."""
        return self._catalog

    @property
    def asset_id(self) -> str:
        return self._asset_id

    @property
    def collection_id(self) -> str:
        return self._collection_id

    @property
    def project_id(self) -> str:
        return self._project_id

class CatalogReadSession(CatalogScopedEntity):
    """Brokered read session entitlement scope; snapshot-bound on server."""

    __slots__ = ("session_id",)

    def __init__(
        self,
        *,
        catalog: "CatalogAPI",
        session_id: str,
        asset_id: str,
        collection_id: str,
        project_id: str,
    ) -> None:
        super().__init__(
            catalog=catalog,
            asset_id=asset_id,
            collection_id=collection_id,
            project_id=project_id,
        )
        self.session_id = session_id

    async def start_run(
        self,
        *,
        batch_size: int = 4096,
        shuffle: dict[str, Any] | None = None,
        shard: dict[str, Any] | CatalogReadShard | None = None,
    ) -> CatalogReadRun:
        control = self._catalog._control_client(project_id=self.project_id)
        shard_policy = shard.as_policy() if isinstance(shard, CatalogReadShard) else shard
        run = await control.create_run(
            self.session_id,
            batch_size=batch_size,
            shuffle_policy=shuffle,
            shard_policy=shard_policy,
        )
        descriptor = await control.get_delivery_descriptor(self.session_id, run["run_id"])
        return CatalogReadRun(session=self, run_id=run["run_id"], delivery=descriptor)

@dataclass(frozen=True)
class CatalogAssetRef:
    """Logical catalog asset identity for brokered reads.

    Not a generic ``Asset`` — catalog-layer identity token only. For product reads,
    hydrate via :class:`~miura.nexus.catalog.catalog_asset.CatalogAsset` or
    ``collection.catalog_asset()``. Batch iteration on the low-level path belongs
    on ``AsyncNexus.catalog``, not ``AsyncCollection``.
    """

    asset_id: str
    collection_id: str
    project_id: str

    def __post_init__(self) -> None:
        if not str(self.asset_id).strip():
            raise ValueError("asset_id is required for CatalogAssetRef")
        if not str(self.collection_id).strip():
            raise ValueError("collection_id is required for CatalogAssetRef")
        if not str(self.project_id).strip():
            raise ValueError("project_id is required for CatalogAssetRef")

    def as_kwargs(self) -> dict[str, str]:
        """Kwargs for ``CatalogAPI.iter_batches`` / ``create_read_session``."""
        return {
            "asset_id": self.asset_id,
            "collection_id": self.collection_id,
            "project_id": self.project_id,
        }

@dataclass(frozen=True)
class CatalogAssetSnapshotRef:
    """Snapshot-bound catalog identity.

    Extends logical ``CatalogAssetRef`` with optional registration/snapshot fields.
    ``as_kwargs()`` passes the logical scope catalog-read-service accepts today;
    snapshot binding is enforced server-side once the bridge is wired.
    """

    asset_id: str
    collection_id: str
    project_id: str
    fileset_snapshot_id: str | None = None
    catalog_registration_id: str | None = None

    def __post_init__(self) -> None:
        CatalogAssetRef(self.asset_id, self.collection_id, self.project_id)

    @property
    def logical_ref(self) -> CatalogAssetRef:
        return CatalogAssetRef(
            asset_id=self.asset_id,
            collection_id=self.collection_id,
            project_id=self.project_id,
        )

    def as_kwargs(self) -> dict[str, str]:
        return self.logical_ref.as_kwargs()

@dataclass(frozen=True)
class CatalogReadShard:
    """Client-side value object: this consumer's catalog-read shard.

    Epico API surface: pass as ``shard=`` on ``catalog.iter_batches`` or
    ``session.start_run``. Users should **not** write ``shard_policy=`` at the
    high level unless using an advanced raw control-plane mode.

    Wire / plan field: ``CatalogReadShard.as_policy()`` serializes to
    ``shard_policy`` on run create and is frozen into ``plan_snapshot``.

    One shard → one catalog-read plan → one run → one delivery ticket → one
    Flight stream.
    """

    strategy: CatalogReadShardStrategy = "rank"
    rank: int | None = None
    world_size: int | None = None
    key: str | None = None

    def __post_init__(self) -> None:
        if self.strategy == "none":
            if self.rank is not None or self.world_size is not None:
                raise ValueError("rank and world_size must be omitted when strategy is none")
            if self.key is not None:
                raise ValueError("key must be omitted when strategy is none")
            return

        if self.rank is None or self.world_size is None:
            raise ValueError("rank and world_size are required when strategy is not none")
        if self.world_size < 1:
            raise ValueError("shard.world_size must be >= 1")
        if self.rank < 0 or self.rank >= self.world_size:
            raise ValueError("shard.rank must satisfy 0 <= rank < world_size")

        if self.strategy == "rank" and self.key is not None:
            raise ValueError("key must be omitted for strategy rank")
        if self.strategy == "rank_partition" and self.key is not None:
            raise ValueError("key must be omitted for strategy rank_partition")

    def as_policy(self) -> dict[str, Any]:
        """Control-plane ``shard_policy`` body for run create (M3+ explicit strategy)."""
        if self.strategy == "none":
            return {"strategy": "none"}
        policy: dict[str, Any] = {
            "strategy": self.strategy,
            "rank": self.rank,
            "world_size": self.world_size,
        }
        if self.key is not None:
            policy["key"] = self.key
        return policy