Source code for ursa.store.factory

"""Config loading + `get_store()` factory.

Search order for the YAML config:

1. Explicit `path` argument
2. `URSA_CONFIG` env var
3. `./ursa.yaml` in the current working directory
4. Packaged `ursa.default.yaml`

The packaged config is profile-agnostic; bucket and credential routing
based on `CONSTELLATION_PROFILE` happens inside constellation-utils.
"""

from __future__ import annotations

import os
from importlib import resources
from pathlib import Path
from typing import Any

import yaml

from ursa.store.backends.local import build_local_store
from ursa.store.backends.r2 import build_r2_store
from ursa.store.base import ObjectStore
from ursa.store.config import (
    LocalStoreConfig,
    R2StoreConfig,
    UrsaConfig,
)

__all__ = [
    "load_config",
    "get_store",
    "ConfigNotFoundError",
]


_DEFAULT_ROLE = "default"
_CONFIG_ENV_VAR = "URSA_CONFIG"
_CWD_FILENAME = "ursa.yaml"
_PACKAGED_CONFIG = "ursa.default.yaml"




class ConfigNotFoundError(FileNotFoundError):
    """Raised when an explicit config path argument or `URSA_CONFIG` env
    points at a file that doesn't exist. (The packaged-default fallback
    is always available, so this only fires when the caller explicitly
    asked for a path that wasn't found.)"""





def _resolve_config_source(path: Path | None) -> tuple[str, Any]:
    """Resolve the config source.

    Returns ``(label, file_obj_or_path)`` where the file contents will
    be parsed as YAML. ``label`` is used in error messages.
    """
    if path is not None:
        if not path.is_file():
            raise ConfigNotFoundError(f"explicit config path not found: {path}")
        return (str(path), path)

    env = os.environ.get(_CONFIG_ENV_VAR)
    if env:
        env_path = Path(env)
        if not env_path.is_file():
            raise ConfigNotFoundError(
                f"${_CONFIG_ENV_VAR} points at {env_path} which does not exist"
            )
        return (str(env_path), env_path)

    cwd_path = Path.cwd() / _CWD_FILENAME
    if cwd_path.is_file():
        return (str(cwd_path), cwd_path)

    return (
        f"<packaged {_PACKAGED_CONFIG}>",
        resources.files("ursa.config") / _PACKAGED_CONFIG,
    )





def load_config(path: Path | None = None) -> UrsaConfig:
    """Load and validate the Ursa config from a YAML file."""
    label, source = _resolve_config_source(path)
    if isinstance(source, Path):
        with source.open("r") as f:
            raw: Any = yaml.safe_load(f)
    else:
        # importlib.resources.abc.Traversable — supports .open
        with source.open("r") as f:
            raw = yaml.safe_load(f)
    if not isinstance(raw, dict):
        raise ValueError(f"{label} did not parse to a mapping (got {type(raw).__name__})")
    return UrsaConfig.model_validate(raw)





def get_store(
    role: str = _DEFAULT_ROLE,
    *,
    config: UrsaConfig | None = None,
    config_path: Path | None = None,
) -> ObjectStore:
    """Return the `ObjectStore` bound to `role` in the active config.

    Pass an already-loaded `config` to avoid re-reading the YAML, or
    `config_path` to override the search order. Without either, the
    YAML is resolved through the search-order described in the module
    docstring.

    Raises `KeyError` (with the available role list) if `role` is not
    defined in the active config.
    """
    cfg = config if config is not None else load_config(config_path)
    try:
        store_cfg = cfg.stores[role]
    except KeyError:
        available = sorted(cfg.stores.keys())
        raise KeyError(
            f"no store configured for role {role!r}; available roles: {available}"
        ) from None
    return _build(store_cfg, role=role)





def _build(
    cfg: R2StoreConfig | LocalStoreConfig,
    *,
    role: str,
) -> ObjectStore:
    if isinstance(cfg, R2StoreConfig):
        return build_r2_store(cfg, role=role)
    if isinstance(cfg, LocalStoreConfig):
        return build_local_store(cfg, role=role)
    # Pydantic's discriminated union should make this unreachable;
    # guard anyway so future additions don't silently fall through.
    raise NotImplementedError(f"unhandled store backend config: {type(cfg).__name__}")