ftm_lakehouse.lake

class Catalog(Generic[DM]):
    """
    Multi-dataset lakehouse catalog.

    The Catalog manages multiple datasets within a lakehouse storage location.

    Example:
        ```python
        from ftm_lakehouse import get_lakehouse

        catalog = get_lakehouse()

        # List datasets
        for dataset in catalog.list_datasets():
            print(dataset.name)

        # Get a specific dataset
        dataset = catalog.get_dataset("my_dataset")
        ```
    """

    def __init__(
        self,
        uri: Uri,
        model_class: type[DM] = DatasetModel,
    ) -> None:
        self.uri = uri
        self._model_class = model_class
        self._log = get_logger(__name__, catalog=mask_uri(uri))

    def __repr__(self) -> str:
        return f"Catalog({mask_uri(self.uri)!r})"

    # -------------------------------------------------------------------------
    # Storage primitives
    # -------------------------------------------------------------------------

    @cached_property
    def _store(self):
        """Raw storage access."""
        return get_store(uri=self.uri, serialization_mode="raw")

    @cached_property
    def _tags(self) -> TagStore:
        """Tag store for freshness tracking."""
        return TagStore(self.uri)

    @cached_property
    def _versions(self) -> VersionStore:
        """Version store for snapshots."""
        return VersionStore(self.uri)

    # -------------------------------------------------------------------------
    # Model access
    # -------------------------------------------------------------------------

    def _load_model(self, **data: Any) -> CatalogModel:
        """Load catalog model from config."""
        return CatalogModel(**load_config(self._store, **data))

    @property
    def model(self) -> CatalogModel:
        """Load and return the catalog model from config.yml."""
        return self._load_model()

    def update_model(self, **data: Any) -> CatalogModel:
        """
        Update config.yml with new data.

        Args:
            **data: Fields to update in the model

        Returns:
            Updated model
        """
        model = self._load_model(**data)
        self._versions.make(path.CONFIG, model)
        return model

    # -------------------------------------------------------------------------
    # Dataset management
    # -------------------------------------------------------------------------

    def get_dataset(self, name: str, **data: Any) -> "Dataset[DM]":
        """
        Get a Dataset instance by name.

        Args:
            name: Dataset name. Validated against
                :func:`ftm_lakehouse.util.validate_dataset_name`.
            **data: Additional config data (auto-saved to config.yml if dataset exists)

        Returns:
            Dataset instance

        Raises:
            ValueError: If ``name`` is not a valid dataset name.
        """
        from ftm_lakehouse.dataset import Dataset

        validate_dataset_name(name)
        dataset_uri = join_uri(self.uri, name)
        dataset = Dataset(
            name=name,
            uri=dataset_uri,
            model_class=self._model_class,
        )

        # Auto-save config if data provided and dataset exists
        if data and dataset.exists():
            dataset.update_model(**data)

        return dataset

    def list_datasets(self) -> Generator["Dataset[DM]", None, None]:
        """
        Iterate through all datasets in the catalog.

        Yields:
            Dataset instances that have a config.yml
        """
        for child in self._store._fs.ls(self.uri):
            dataset_name = Path(child).name
            if self._store.exists(f"{dataset_name}/{path.CONFIG}"):
                yield self.get_dataset(dataset_name)

    def create_dataset(self, name: str, **data: Any) -> "Dataset[DM]":
        """
        Create a new dataset.

        Args:
            name: Dataset name
            **data: Initial config data, recorded into ``config.yml`` at
                creation (e.g. ``shards=8`` for a huge dataset)

        Returns:
            Created Dataset instance
        """
        dataset = self.get_dataset(name, **data)
        dataset.ensure(**data)
        return dataset

class Dataset(Generic[DM]):
    """
    A single dataset within the lakehouse – a stateless handle around the
    dataset's name, storage uri and configuration.

    Repository access goes through the LRU-cached factories, so every path
    addressing the same dataset – this handle, the module-level
    ``get_entities("name")`` convenience functions, and operations – shares
    one repository instance:

    - `get_archive()`: File storage (ArchiveRepository)
    - `get_entities()`: Entity/statement operations (EntityRepository)
    - `get_documents()`: Document metadata (DocumentRepository)
    - `get_mappings()`: Mapping configurations (MappingRepository)

    Example:
        ```python
        from ftm_lakehouse import get_dataset

        dataset = get_dataset("my_dataset")
        dataset.ensure()

        # Add entities
        dataset.get_entities().add(entity, origin="import")

        # Archive files
        dataset.get_archive().store(uri)

        # Update config
        dataset.update_model(title="New Title")
        ```
    """

    def __init__(
        self,
        name: str,
        uri: Uri,
        model_class: type[DM] = DatasetModel,
    ) -> None:
        self.name = name
        self.uri = uri
        self._model_class = model_class
        self._settings = Settings()
        self._log = log.bind(dataset=name, uri=mask_uri(uri))

        if self._store.is_local and self._settings.on_zfs:
            if self._settings.zfs_pool is None:
                raise RuntimeError("Configure LAKEHOUSE_ZFS_POOL for zfs integration!")
            ensure_zfs_dataset(self._settings.zfs_pool, self.name)

    def __repr__(self) -> str:
        return f"Dataset({self.name!r})"

    # -------------------------------------------------------------------------
    # Storage primitives
    # -------------------------------------------------------------------------

    @property
    def _store(self) -> Store:
        """Raw storage access."""
        return get_store(uri=ensure_api_uri(self.uri), serialization_mode="raw")

    @property
    def _versions(self) -> VersionStore:
        """Version store for snapshots."""
        return get_versions(self.name, self.uri)

    # -------------------------------------------------------------------------
    # Model access (config.yml via VersionStore)
    # -------------------------------------------------------------------------

    def _load_model(self, **data: Any) -> DM:
        """Load dataset model from config.yml."""
        data["name"] = self.name
        data.pop("storage", None)
        return self._model_class(**load_config(self._store, **data))

    @property
    def model(self) -> DM:
        """Load and return the dataset model from config.yml."""
        return self._load_model()

    @property
    def index(self) -> DM:
        """Load and return the generated index.json (or fall back to config.yml)"""
        index = self._versions.get(path.INDEX, model=self._model_class)
        if index:
            return index
        return self.model

    def update_model(self, **data: Any) -> DM:
        """
        Update config.yml with new data.

        Uses VersionStore to create versioned snapshots.

        Args:
            **data: Fields to update in the model

        Returns:
            Updated model
        """
        model = self._load_model(**data)
        self._versions.make(path.CONFIG, model)
        self._log.info("Updated dataset config")
        return model

    # -------------------------------------------------------------------------
    # Repositories (resolved through the LRU-cached factories)
    # -------------------------------------------------------------------------

    def get_archive(self) -> ArchiveRepository:
        """File archive operations."""
        return get_archive(self.name, self.uri)

    def get_entities(self) -> EntityRepository:
        """Entity/statement operations."""
        return get_entities(self.name, self.uri)

    def get_documents(self) -> DocumentRepository:
        """Document metadata operations."""
        return get_documents(self.name, self.uri)

    def get_mappings(self) -> MappingRepository:
        """Mapping configuration storage."""
        return get_mappings(self.name, self.uri)

    # -------------------------------------------------------------------------
    # Lifecycle
    # -------------------------------------------------------------------------

    def exists(self) -> bool:
        """Check if dataset exists (has config.yml)."""
        return self._store.exists(path.CONFIG)

    def ensure(self, **data: Any) -> None:
        """Ensure dataset exists, create config.yml if needed.

        Args:
            **data: Initial config data recorded at creation (e.g.
                ``shards=8`` for a huge dataset). Ignored when the dataset
                already exists.
        """
        if not self.exists():
            self.update_model(**data)
            self._log.info("Created dataset")