metaxy 0.0.0__py3-none-any.whl

metaxy/cli/push.py

@@ -0,0 +1,55 @@
+"""Push command for recording feature versions."""
+
+from rich.console import Console
+
+console = Console()
+
+
+def push(store: str | None = None):
+    """Record all feature versions (push graph snapshot).
+
+    Records all features in the active graph to the metadata store
+    with a deterministic snapshot version. This should be run after deploying
+    new feature definitions.
+
+    Example:
+        $ metaxy push
+
+        ✓ Recorded feature graph
+          Snapshot version: abc123def456...
+
+        # Or if already recorded:
+        ℹ Snapshot already recorded (skipped)
+          Snapshot version: abc123def456...
+
+    Args:
+        store: The metadata store to use. Defaults to the default store.
+    """
+    from metaxy.cli.context import get_store
+    from metaxy.models.feature import FeatureGraph
+
+    metadata_store = get_store(store)
+
+    with metadata_store:
+        # Get active graph
+        active_graph = FeatureGraph.get_active()
+        if len(active_graph.features_by_key) == 0:
+            console.print("[yellow]⚠[/yellow] No features in active graph")
+            return
+
+        # Record feature graph snapshot (idempotent)
+        # Returns (snapshot_version, already_exists)
+        snapshot_version, already_exists = (
+            metadata_store.record_feature_graph_snapshot()
+        )
+
+        if already_exists:
+            console.print("[blue]ℹ[/blue] Snapshot already recorded (skipped)")
+            console.print(f"  Snapshot version: {snapshot_version}")
+        else:
+            console.print("[green]✓[/green] Recorded feature graph")
+            console.print(f"  Snapshot version: {snapshot_version}")
+
+
+if __name__ == "__main__":
+    push()

metaxy/config.py

@@ -0,0 +1,450 @@
+"""Configuration system for Metaxy using pydantic-settings."""
+# pyright: reportImportCycles=false
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, TypeVar
+
+try:
+    import tomllib  # Python 3.11+  # pyright: ignore[reportMissingImports]
+except ImportError:
+    import tomli as tomllib  # Fallback for Python 3.10
+
+import warnings
+from contextvars import ContextVar
+
+from pydantic import Field as PydanticField
+from pydantic import PrivateAttr
+from pydantic_settings import (
+    BaseSettings,
+    PydanticBaseSettingsSource,
+    SettingsConfigDict,
+)
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from metaxy.metadata_store.base import (
+        MetadataStore,  # pyright: ignore[reportImportCycles]
+    )
+
+T = TypeVar("T")
+
+
+class TomlConfigSettingsSource(PydanticBaseSettingsSource):
+    """Custom settings source for TOML configuration files.
+
+    Auto-discovers configuration in this order:
+    1. Explicit file path if provided
+    2. metaxy.toml in current directory (preferred)
+    3. pyproject.toml [tool.metaxy] section (fallback)
+    4. No config (returns empty dict)
+    """
+
+    def __init__(self, settings_cls: type[BaseSettings], toml_file: Path | None = None):
+        super().__init__(settings_cls)
+        self.toml_file = toml_file or self._discover_config_file()
+        self.toml_data = self._load_toml()
+
+    def _discover_config_file(self) -> Path | None:
+        """Auto-discover config file."""
+        # Prefer metaxy.toml
+        if Path("metaxy.toml").exists():
+            return Path("metaxy.toml")
+
+        # Fallback to pyproject.toml
+        if Path("pyproject.toml").exists():
+            return Path("pyproject.toml")
+
+        return None
+
+    def _load_toml(self) -> dict[str, Any]:
+        """Load TOML file and extract metaxy config."""
+        if self.toml_file is None:
+            return {}
+
+        with open(self.toml_file, "rb") as f:
+            data = tomllib.load(f)
+
+        # Extract [tool.metaxy] from pyproject.toml or root from metaxy.toml
+        if self.toml_file.name == "pyproject.toml":
+            return data.get("tool", {}).get("metaxy", {})
+        else:
+            return data
+
+    def get_field_value(self, field: Any, field_name: str) -> tuple[Any, str, bool]:
+        """Get field value from TOML data."""
+        field_value = self.toml_data.get(field_name)
+        return field_value, field_name, False
+
+    def __call__(self) -> dict[str, Any]:
+        """Return all settings from TOML."""
+        return self.toml_data
+
+
+class StoreConfig(BaseSettings):
+    """Configuration for a single metadata store.
+
+    Structure:
+        type: Full import path to store class
+        config: Dict of all configuration (including fallback_stores)
+
+    Example:
+        >>> config = StoreConfig(
+        ...     type="metaxy_delta.DeltaMetadataStore",
+        ...     config={
+        ...         "table_uri": "s3://bucket/metadata",
+        ...         "region": "us-west-2",
+        ...         "fallback_stores": ["prod"],
+        ...     }
+        ... )
+    """
+
+    model_config = SettingsConfigDict(
+        extra="forbid",  # Only type and config fields allowed
+    )
+
+    # Store class (full import path)
+    type: str
+
+    # Store configuration (all kwargs for __init__)
+    # This includes fallback_stores, table_uri, db_path, storage_options, etc.
+    config: dict[str, Any] = PydanticField(default_factory=dict)
+
+
+class PluginConfig(BaseSettings):
+    """Configuration for Metaxy plugins"""
+
+    enabled: bool = PydanticField(
+        default=False,
+        description="Whether to enable plugin.",
+    )
+
+    _plugin: str = PrivateAttr()
+
+
+class SQLModelConfig(PluginConfig):
+    """Configuration for SQLModel"""
+
+    infer_db_table_names: bool = PydanticField(
+        default=True,
+        description="Whether to automatically use `FeatureKey.table_name` for sqlalchemy's __tablename__ value.",
+    )
+
+    # Whether to use SQLModel definitions for system tables (for Alembic migrations)
+    system_tables: bool = PydanticField(
+        default=False,
+        description="Whether to use SQLModel definitions for system tables (for Alembic migrations).",
+    )
+
+    _plugin: str = PrivateAttr(default="sqlmodel")
+
+
+class ExtConfig(BaseSettings):
+    """Configuration for Metaxy integrations with third-party tools"""
+
+    model_config = SettingsConfigDict(
+        extra="allow",
+    )
+
+    sqlmodel: SQLModelConfig = PydanticField(default_factory=SQLModelConfig)
+
+
+# Context variable for storing the app context
+_metaxy_config: ContextVar["MetaxyConfig | None"] = ContextVar(
+    "_metaxy_config", default=None
+)
+
+
+class MetaxyConfig(BaseSettings):
+    """Main Metaxy configuration.
+
+    Loads from:
+    1. TOML file (metaxy.toml or pyproject.toml [tool.metaxy])
+    2. Environment variables (METAXY_*)
+    3. Init arguments
+
+    Priority: init > env vars > TOML
+
+    Example:
+        >>> # Auto-discover config
+        >>> config = MetaxyConfig.load()
+        >>>
+        >>> # Get store instance
+        >>> store = config.get_store("prod")
+        >>>
+        >>> # Override via env var
+        >>> # METAXY_STORE=staging METAXY_REGISTRY=myapp.features:my_graph
+        >>> config = MetaxyConfig.load()
+        >>> store = config.get_store()  # Uses staging with custom graph
+    """
+
+    model_config = SettingsConfigDict(
+        env_prefix="METAXY_",
+        env_nested_delimiter="__",
+    )
+
+    # Store to use
+    store: str = "dev"
+
+    # Named store configurations
+    stores: dict[str, StoreConfig] = PydanticField(default_factory=dict)
+
+    # Migrations directory
+    migrations_dir: str = ".metaxy/migrations"
+
+    # Entrypoints to load (list of module paths)
+    entrypoints: list[str] = PydanticField(default_factory=list)
+
+    # Graph rendering theme
+    theme: str = "default"
+
+    ext: ExtConfig = PydanticField(default_factory=ExtConfig)
+
+    @property
+    def plugins(self) -> list[str]:
+        """Returns all enabled plugin names from ext configuration."""
+        plugins = []
+        for field_name in type(self.ext).model_fields:
+            field_value = getattr(self.ext, field_name)
+            if hasattr(field_value, "_plugin") and field_value.enabled:
+                plugins.append(field_value._plugin)
+        return plugins
+
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        """Customize settings sources: init → env → TOML.
+
+        Priority (first wins):
+        1. Init arguments
+        2. Environment variables
+        3. TOML file
+        """
+        toml_settings = TomlConfigSettingsSource(settings_cls)
+        return (init_settings, env_settings, toml_settings)
+
+    @classmethod
+    def get(cls) -> "MetaxyConfig":
+        """Get the current Metaxy configuration."""
+        cfg = _metaxy_config.get()
+        if cfg is None:
+            warnings.warn(
+                UserWarning(
+                    "Global Metaxy configuration not initialized. It can be set with MetaxyConfig.set(config) typically after loading it from a toml file. Returning default configuration (with environment variables and other pydantic settings sources resolved)."
+                )
+            )
+            return cls()
+        else:
+            return cfg
+
+    @classmethod
+    def set(cls, config: Self) -> None:
+        """Set the current Metaxy configuration."""
+        _metaxy_config.set(config)
+
+    @classmethod
+    def load(
+        cls, config_file: str | Path | None = None, *, search_parents: bool = True
+    ) -> "MetaxyConfig":
+        """Load config with auto-discovery and parent directory search.
+
+        Args:
+            config_file: Optional config file path (overrides auto-discovery)
+            search_parents: Search parent directories for config file (default: True)
+
+        Returns:
+            Loaded config (TOML + env vars merged)
+
+        Example:
+            >>> # Auto-discover with parent search
+            >>> config = MetaxyConfig.load()
+            >>>
+            >>> # Explicit file
+            >>> config = MetaxyConfig.load("custom.toml")
+
+            >>> # Auto-discover without parent search
+            >>> config = MetaxyConfig.load(search_parents=False)
+        """
+        # Search for config file if not explicitly provided
+        if config_file is None and search_parents:
+            config_file = cls._discover_config_with_parents()
+
+        # For explicit file, temporarily patch the TomlConfigSettingsSource
+        # to use that file, then use normal instantiation
+        # This ensures env vars still work
+
+        if config_file:
+            # Create a custom settings source class for this file
+            toml_path = Path(config_file)
+
+            class CustomTomlSource(TomlConfigSettingsSource):
+                def __init__(self, settings_cls: type[BaseSettings]):
+                    # Skip auto-discovery, use explicit file
+                    super(TomlConfigSettingsSource, self).__init__(settings_cls)
+                    self.toml_file = toml_path
+                    self.toml_data = self._load_toml()
+
+            # Customize sources to use custom TOML file
+            original_method = cls.settings_customise_sources
+
+            @classmethod  # type: ignore[misc]
+            def custom_sources(
+                cls_inner,
+                settings_cls,
+                init_settings,
+                env_settings,
+                dotenv_settings,
+                file_secret_settings,
+            ):
+                toml_settings = CustomTomlSource(settings_cls)
+                return (init_settings, env_settings, toml_settings)
+
+            # Temporarily replace method
+            cls.settings_customise_sources = custom_sources  # type: ignore[assignment]
+            config = cls()
+            cls.settings_customise_sources = original_method  # type: ignore[method-assign]
+        else:
+            # Use default sources (auto-discovery + env vars)
+            config = cls()
+
+        cls.set(config)
+
+        return config
+
+    @staticmethod
+    def _discover_config_with_parents() -> Path | None:
+        """Discover config file by searching current and parent directories.
+
+        Searches for metaxy.toml or pyproject.toml in current directory,
+        then iteratively searches parent directories.
+
+        Returns:
+            Path to config file if found, None otherwise
+        """
+        current = Path.cwd()
+
+        while True:
+            # Check for metaxy.toml (preferred)
+            metaxy_toml = current / "metaxy.toml"
+            if metaxy_toml.exists():
+                return metaxy_toml
+
+            # Check for pyproject.toml
+            pyproject_toml = current / "pyproject.toml"
+            if pyproject_toml.exists():
+                return pyproject_toml
+
+            # Move to parent
+            parent = current.parent
+            if parent == current:
+                # Reached root
+                break
+            current = parent
+
+        return None
+
+    def get_store(
+        self,
+        name: str | None = None,
+    ) -> "MetadataStore":
+        """Instantiate metadata store by name.
+
+        Args:
+            name: Store name (uses config.store if None)
+
+        Returns:
+            Instantiated metadata store
+
+        Raises:
+            ValueError: If store name not found in config, or if fallback stores
+                have different hash algorithms than the parent store
+            ImportError: If store class cannot be imported
+
+        Example:
+            >>> config = MetaxyConfig.load()
+            >>> store = config.get_store("prod")
+            >>>
+            >>> # Use default store
+            >>> store = config.get_store()
+        """
+        from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+
+        if len(self.stores) == 0:
+            raise ValueError(
+                "No Metaxy stores available. They should be configured in metaxy.toml|pyproject.toml or via environment variables."
+            )
+
+        name = name or self.store
+
+        if name not in self.stores:
+            raise ValueError(
+                f"Store '{name}' not found in config. "
+                f"Available stores: {list(self.stores.keys())}"
+            )
+
+        store_config = self.stores[name]
+
+        # Import store class
+        store_class = self._import_class(store_config.type)
+
+        # Extract configuration
+        config_copy = store_config.config.copy()
+        fallback_store_names = config_copy.pop("fallback_stores", [])
+
+        # Get hash_algorithm from config (if specified) and convert to enum
+        configured_hash_algorithm = config_copy.get("hash_algorithm")
+        if configured_hash_algorithm is not None:
+            # Convert string to enum if needed
+            if isinstance(configured_hash_algorithm, str):
+                configured_hash_algorithm = HashAlgorithm(configured_hash_algorithm)
+                config_copy["hash_algorithm"] = configured_hash_algorithm
+        else:
+            # Use default
+            configured_hash_algorithm = HashAlgorithm.XXHASH64
+            config_copy["hash_algorithm"] = configured_hash_algorithm
+
+        # Build fallback stores recursively
+        fallback_stores = []
+        for fallback_name in fallback_store_names:
+            fallback_store = self.get_store(fallback_name)
+            fallback_stores.append(fallback_store)
+
+        # Instantiate store with config + fallback_stores
+        store = store_class(
+            fallback_stores=fallback_stores,
+            **config_copy,
+        )
+
+        # Verify the store actually uses the hash algorithm we configured
+        # (in case a store subclass overrides the default or ignores the parameter)
+        if store.hash_algorithm != configured_hash_algorithm:
+            raise ValueError(
+                f"Store '{name}' ({store_class.__name__}) was configured with "
+                f"hash_algorithm='{configured_hash_algorithm.value}' but is using "
+                f"'{store.hash_algorithm.value}'. The store class may have overridden "
+                f"the hash algorithm. All stores must use the same hash algorithm."
+            )
+
+        return store
+
+    @staticmethod
+    def _import_class(class_path: str) -> type:
+        """Import class from module path.
+
+        Args:
+            class_path: Full import path like "metaxy.metadata_store.InMemoryMetadataStore"
+
+        Returns:
+            Imported class
+
+        Raises:
+            ImportError: If module or class not found
+        """
+        module_path, class_name = class_path.rsplit(".", 1)
+        module = __import__(module_path, fromlist=[class_name])
+        return getattr(module, class_name)

metaxy/data_versioning/__init__.py

@@ -0,0 +1,24 @@
+"""Data versioning module for sample-level data version calculation."""
+
+from metaxy.data_versioning.calculators import (
+    DataVersionCalculator,
+    PolarsDataVersionCalculator,
+)
+from metaxy.data_versioning.diff import (
+    DiffResult,
+    MetadataDiffResolver,
+    NarwhalsDiffResolver,
+)
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+from metaxy.data_versioning.joiners import NarwhalsJoiner, UpstreamJoiner
+
+__all__ = [
+    "HashAlgorithm",
+    "UpstreamJoiner",
+    "NarwhalsJoiner",
+    "DataVersionCalculator",
+    "PolarsDataVersionCalculator",
+    "DiffResult",
+    "MetadataDiffResolver",
+    "NarwhalsDiffResolver",
+]

metaxy/data_versioning/calculators/__init__.py

@@ -0,0 +1,13 @@
+"""Data version calculators for computing hash from upstream data."""
+
+from metaxy.data_versioning.calculators.base import DataVersionCalculator
+from metaxy.data_versioning.calculators.duckdb import DuckDBDataVersionCalculator
+from metaxy.data_versioning.calculators.ibis import IbisDataVersionCalculator
+from metaxy.data_versioning.calculators.polars import PolarsDataVersionCalculator
+
+__all__ = [
+    "DataVersionCalculator",
+    "DuckDBDataVersionCalculator",
+    "IbisDataVersionCalculator",
+    "PolarsDataVersionCalculator",
+]

metaxy/data_versioning/calculators/base.py

@@ -0,0 +1,97 @@
+"""Abstract base class for data version calculators."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import narwhals as nw
+
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+
+if TYPE_CHECKING:
+    from metaxy.models.feature_spec import FeatureSpec
+    from metaxy.models.plan import FeaturePlan
+
+
+class DataVersionCalculator(ABC):
+    """Calculates data_version hash from joined upstream data.
+
+    The calculator takes joined upstream data (output from UpstreamJoiner)
+    and computes the data_version hash for each sample.
+
+    This is Step 2 in the data versioning process:
+    1. Join upstream features → unified upstream view
+    2. Calculate data_version from upstream → target versions ← THIS STEP
+    3. Diff with current metadata → identify changes
+
+    All calculators work with Narwhals LazyFrames for backend compatibility.
+
+    Examples:
+        - PolarsDataVersionCalculator: Uses polars-hash for in-memory hashing
+        - NarwhalsDataVersionCalculator: Uses native SQL hash functions in the database
+    """
+
+    @property
+    @abstractmethod
+    def supported_algorithms(self) -> list[HashAlgorithm]:
+        """List of hash algorithms this calculator supports.
+
+        Returns:
+            List of supported HashAlgorithm values
+
+        Example:
+            >>> calc = PolarsDataVersionCalculator()
+            >>> HashAlgorithm.XXHASH64 in calc.supported_algorithms
+            True
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def default_algorithm(self) -> HashAlgorithm:
+        """Default hash algorithm for this calculator.
+
+        Should be the most performant algorithm that's widely compatible.
+        Typically xxHash64 for cross-database compatibility.
+
+        Returns:
+            Default HashAlgorithm
+        """
+        pass
+
+    @abstractmethod
+    def calculate_data_versions(
+        self,
+        joined_upstream: nw.LazyFrame[Any],
+        feature_spec: "FeatureSpec",
+        feature_plan: "FeaturePlan",
+        upstream_column_mapping: dict[str, str],
+        hash_algorithm: HashAlgorithm | None = None,
+    ) -> nw.LazyFrame[Any]:
+        """Calculate data_version column from joined upstream data.
+
+        Computes a Merkle tree hash for each sample by:
+        1. For each field in the feature:
+           a. Concatenate: field_key | code_version | upstream hashes
+           b. Hash the concatenated string
+        2. Create struct with all field hashes
+        3. Add as data_version column
+
+        Args:
+            joined_upstream: Narwhals LazyFrame with all upstream data_version columns joined
+                (output from UpstreamJoiner.join_upstream)
+            feature_spec: Specification of the feature being computed
+            feature_plan: Resolved feature plan with dependencies
+            upstream_column_mapping: Maps upstream feature key -> column name
+                where its data_version struct is located in joined_upstream
+                Example: {"video": "__upstream_video__data_version"}
+            hash_algorithm: Hash algorithm to use. If None, uses self.default_algorithm.
+                Must be in self.supported_algorithms.
+
+        Returns:
+            Narwhals LazyFrame with data_version column added
+            Shape: [sample_uid, __upstream_*__data_version columns, data_version (new)]
+
+        Raises:
+            ValueError: If hash_algorithm not in supported_algorithms
+        """
+        pass