metaxy 0.0.0__py3-none-any.whl

metaxy/metadata_store/sqlite.py

@@ -0,0 +1,187 @@
+"""SQLite metadata store - thin wrapper around IbisMetadataStore."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+if TYPE_CHECKING:
+    from metaxy.metadata_store.base import MetadataStore
+
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+from metaxy.metadata_store.ibis import IbisMetadataStore
+
+
+class SQLiteMetadataStore(IbisMetadataStore):
+    """
+    SQLite metadata store using Ibis backend.
+
+    Convenience wrapper that configures IbisMetadataStore for SQLite.
+
+    Hash algorithm support:
+    - MD5: Available (built-in SQLite function via extension)
+
+    Components:
+        - joiner: NarwhalsJoiner (works with any backend)
+        - calculator: PolarsDataVersionCalculator (SQLite always uses Polars, no native compute)
+        - diff_resolver: NarwhalsDiffResolver
+
+    Examples:
+        >>> # Local file database
+        >>> with SQLiteMetadataStore("metadata.db") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # In-memory database
+        >>> with SQLiteMetadataStore(":memory:") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # Explicit path
+        >>> store = SQLiteMetadataStore(Path("/path/to/metadata.db"))
+        >>> with store:
+        ...     store.write_metadata(MyFeature, df)
+    """
+
+    def __init__(
+        self,
+        database: str | Path,
+        *,
+        fallback_stores: list["MetadataStore"] | None = None,
+        **kwargs,
+    ):
+        """
+        Initialize SQLite metadata store.
+
+        Args:
+            database: Database connection string or path.
+                - File path: "metadata.db" or Path("metadata.db")
+                - In-memory: ":memory:"
+
+                Note: Parent directories are NOT created automatically. Ensure paths exist
+                before initializing the store.
+            fallback_stores: Ordered list of read-only fallback stores.
+            **kwargs: Passed to IbisMetadataStore (e.g., hash_algorithm, prefer_native)
+        """
+        database_str = str(database)
+
+        # Build connection params for Ibis SQLite backend
+        connection_params = {"database": database_str}
+
+        self.database = database_str
+
+        # Initialize Ibis store with SQLite backend
+        super().__init__(
+            backend="sqlite",
+            connection_params=connection_params,
+            fallback_stores=fallback_stores,
+            **kwargs,
+        )
+
+    def _get_default_hash_algorithm(self) -> HashAlgorithm:
+        """Get default hash algorithm for SQLite stores.
+
+        Uses MD5 which is universally supported in SQLite.
+        """
+        return HashAlgorithm.MD5
+
+    @classmethod
+    def supports_structs(cls) -> bool:
+        """SQLite does not support struct types natively.
+
+        Returns:
+            False - SQLite stores structs as JSON strings
+        """
+        return False
+
+    def _supports_native_components(self) -> bool:
+        """SQLite stores do not support native data version calculations.
+
+        SQLite doesn't have built-in hash functions (MD5, SHA256, etc.),
+        so we always use Polars components for data versioning.
+        """
+        return False
+
+    def _serialize_for_storage(self, df: pl.DataFrame) -> pl.DataFrame:
+        """Serialize structs and arrays to JSON strings for SQLite storage.
+
+        SQLite doesn't support struct or array types, so we convert them to JSON strings.
+
+        Args:
+            df: DataFrame with potential struct/array columns
+
+        Returns:
+            DataFrame with struct/array columns converted to JSON strings
+        """
+        # Convert struct and array columns to JSON strings
+        for col_name in df.columns:
+            dtype = df.schema[col_name]
+            if isinstance(dtype, pl.Struct):
+                # Convert struct to JSON string
+                df = df.with_columns(
+                    pl.col(col_name).struct.json_encode().alias(col_name)
+                )
+            elif isinstance(dtype, pl.List):
+                # Convert array/list to JSON string
+                # Note: Polars doesn't have native list.json_encode(), so we use map_elements
+                import json
+
+                df = df.with_columns(
+                    pl.col(col_name)
+                    .map_elements(
+                        lambda x: None if x is None else json.dumps(x),
+                        return_dtype=pl.Utf8,
+                    )
+                    .alias(col_name)
+                )
+
+        return df
+
+    def _deserialize_from_storage(self, df: pl.DataFrame) -> pl.DataFrame:
+        """Deserialize JSON strings back to structs and arrays.
+
+        Converts JSON string columns back to their original struct/array types.
+
+        Args:
+            df: DataFrame with JSON string columns
+
+        Returns:
+            DataFrame with JSON strings converted back to structs/arrays
+        """
+        # Known struct and array columns with their expected dtypes
+        # data_version is a struct, fields is a list of structs
+        # Migration system columns: operation_ids, expected_steps (list of strings),
+        #                          migration_yaml (struct), affected_features (list of strings)
+
+        # Columns that need JSON deserialization with specific dtypes
+        json_columns = {
+            "data_version": None,  # Infer from data
+            "migration_yaml": None,  # Infer from data
+            # "feature_spec": Leave as JSON string - contains enum values that can't be parsed
+            "operation_ids": pl.List(pl.Utf8),  # List of strings
+            "expected_steps": pl.List(pl.Utf8),  # List of strings
+            "affected_features": pl.List(pl.Utf8),  # List of strings
+        }
+
+        # Deserialize JSON columns
+        for col_name, dtype in json_columns.items():
+            if col_name in df.columns and df.schema[col_name] == pl.Utf8:
+                if len(df) > 0:
+                    if dtype is None:
+                        # Infer dtype from sample value
+                        sample_value = df[col_name].drop_nulls().head(1)
+                        if len(sample_value) > 0:
+                            inferred_series = sample_value.str.json_decode()
+                            inferred_dtype = inferred_series.dtype
+                            df = df.with_columns(
+                                pl.col(col_name)
+                                .str.json_decode(dtype=inferred_dtype)
+                                .alias(col_name)
+                            )
+                    else:
+                        # Use provided dtype
+                        df = df.with_columns(
+                            pl.col(col_name)
+                            .str.json_decode(dtype=dtype)
+                            .alias(col_name)
+                        )
+
+        return df

metaxy/metadata_store/system_tables.py

@@ -0,0 +1,257 @@
+"""System table storage layer for metadata store.
+
+Provides type-safe access to migration system tables using struct-based storage.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from datetime import datetime, timezone
+from typing import Any
+
+import narwhals as nw
+import polars as pl
+
+from metaxy.metadata_store._protocols import MetadataStoreProtocol
+from metaxy.models.types import FeatureKey
+
+# System namespace
+SYSTEM_NAMESPACE = "metaxy-system"
+
+# System table keys
+FEATURE_VERSIONS_KEY = FeatureKey([SYSTEM_NAMESPACE, "feature_versions"])
+MIGRATION_EVENTS_KEY = FeatureKey([SYSTEM_NAMESPACE, "migration_events"])
+# Note: No migrations table - definitions live in YAML files, only events are stored
+
+# Context variable for suppressing feature_version warning in migrations
+_suppress_feature_version_warning: ContextVar[bool] = ContextVar(
+    "_suppress_feature_version_warning", default=False
+)
+
+
+@contextmanager
+def allow_feature_version_override() -> Iterator[None]:
+    """
+    Context manager to suppress warnings when writing metadata with pre-existing feature_version.
+
+    This should only be used in migration code where writing historical feature versions
+    is intentional and necessary.
+
+    Example:
+        >>> with allow_feature_version_override():
+        ...     # DataFrame already has feature_version column from migration
+        ...     store.write_metadata(MyFeature, df_with_feature_version)
+    """
+    token = _suppress_feature_version_warning.set(True)
+    try:
+        yield
+    finally:
+        _suppress_feature_version_warning.reset(token)
+
+
+# Common Polars schemas for system tables
+# TODO: Migrate to use METAXY_*_COL constants instead of plain names
+FEATURE_VERSIONS_SCHEMA = {
+    "feature_key": pl.String,
+    "feature_version": pl.String,  # TODO: Use METAXY_FEATURE_VERSION_COL
+    "recorded_at": pl.Datetime("us"),
+    "feature_spec": pl.String,
+    "feature_class_path": pl.String,
+    "snapshot_version": pl.String,  # TODO: Use METAXY_SNAPSHOT_ID_COL
+}
+
+MIGRATION_EVENTS_SCHEMA = {
+    "migration_id": pl.String,
+    "event_type": pl.String,  # "started", "feature_started", "feature_completed", "completed", "failed"
+    "timestamp": pl.Datetime("us"),
+    "feature_key": pl.String,  # Empty for migration-level events
+    "rows_affected": pl.Int64,
+    "error_message": pl.String,  # Empty if no error
+}
+
+
+class SystemTableStorage:
+    """Storage layer for migration system tables.
+
+    Provides type-safe access to migration snapshots, migrations, and events.
+    Uses struct-based storage (not JSON/bytes) for efficient queries.
+
+    Status is computed at query-time from events (append-only).
+    """
+
+    def __init__(self, store: MetadataStoreProtocol):
+        """Initialize storage layer.
+
+        Args:
+            store: Metadata store to use for system tables
+        """
+        self.store = store
+
+    # ========== Migrations ==========
+    # Note: Migration definitions are stored in YAML files (git), not in the database.
+    # Only execution events are stored in DB for tracking progress and state.
+
+    def list_executed_migrations(self) -> list[str]:
+        """List all migration IDs that have execution events.
+
+        Returns:
+            List of migration IDs that have been started/executed
+        """
+        lazy = self.store._read_metadata_native(MIGRATION_EVENTS_KEY)
+
+        if lazy is None:
+            return []
+
+        df = lazy.select("migration_id").unique().collect().to_polars()
+        return df["migration_id"].to_list()
+
+    # ========== Events ==========
+
+    def write_event(
+        self,
+        migration_id: str,
+        event_type: str,
+        feature_key: str = "",
+        rows_affected: int = 0,
+        error_message: str = "",
+    ) -> None:
+        """Write migration event to system table (append-only).
+
+        Args:
+            migration_id: Migration this event belongs to
+            event_type: Event type ("started", "feature_started", "feature_completed", "completed", "failed")
+            feature_key: Feature key (empty for migration-level events)
+            rows_affected: Number of rows affected (for feature events)
+            error_message: Error message (empty if no error)
+        """
+        record = pl.DataFrame(
+            {
+                "migration_id": [migration_id],
+                "event_type": [event_type],
+                "timestamp": [datetime.now(timezone.utc)],
+                "feature_key": [feature_key],
+                "rows_affected": [rows_affected],
+                "error_message": [error_message],
+            },
+            schema=MIGRATION_EVENTS_SCHEMA,
+        )
+        self.store._write_metadata_impl(MIGRATION_EVENTS_KEY, record)
+
+    def get_migration_events(self, migration_id: str) -> nw.LazyFrame[Any]:
+        """Get all events for a migration.
+
+        Args:
+            migration_id: Migration ID
+
+        Returns:
+            Lazy frame with events sorted by timestamp
+        """
+        lazy = self.store._read_metadata_native(
+            MIGRATION_EVENTS_KEY,
+            filters=[nw.col("migration_id") == migration_id],
+        )
+
+        if lazy is None:
+            # No events yet
+            return nw.from_native(pl.DataFrame(schema=MIGRATION_EVENTS_SCHEMA).lazy())
+
+        return lazy.sort("timestamp", descending=False)
+
+    def get_migration_status(self, migration_id: str) -> str:
+        """Compute migration status from events at query-time.
+
+        Args:
+            migration_id: Migration ID
+
+        Returns:
+            Status: "not_started", "in_progress", "completed", "failed"
+        """
+        events_lazy = self.get_migration_events(migration_id)
+        events_df = events_lazy.collect().to_polars()
+
+        if events_df.height == 0:
+            return "not_started"
+
+        # Get latest event
+        latest_event = events_df.sort("timestamp", descending=True).head(1)
+        latest_event_type = latest_event["event_type"][0]
+
+        if latest_event_type == "completed":
+            return "completed"
+        elif latest_event_type == "failed":
+            return "failed"
+        elif latest_event_type in ("started", "feature_started", "feature_completed"):
+            return "in_progress"
+
+        return "not_started"
+
+    def is_feature_completed(self, migration_id: str, feature_key: str) -> bool:
+        """Check if a specific feature completed successfully in a migration.
+
+        Args:
+            migration_id: Migration ID
+            feature_key: Feature key to check
+
+        Returns:
+            True if feature completed without errors
+        """
+        events_lazy = self.get_migration_events(migration_id)
+        events_df = (
+            events_lazy.filter(
+                (nw.col("feature_key") == feature_key)
+                & (nw.col("event_type") == "feature_completed")
+                & (nw.col("error_message") == "")
+            )
+            .collect()
+            .to_polars()
+        )
+
+        return events_df.height > 0
+
+    def get_completed_features(self, migration_id: str) -> list[str]:
+        """Get list of features that completed successfully in a migration.
+
+        Args:
+            migration_id: Migration ID
+
+        Returns:
+            List of feature keys
+        """
+        events_lazy = self.get_migration_events(migration_id)
+        events_df = (
+            events_lazy.filter(
+                (nw.col("event_type") == "feature_completed")
+                & (nw.col("error_message") == "")
+            )
+            .collect()
+            .to_polars()
+        )
+
+        return events_df["feature_key"].unique().to_list()
+
+    def get_failed_features(self, migration_id: str) -> dict[str, str]:
+        """Get features that failed in a migration with error messages.
+
+        Args:
+            migration_id: Migration ID
+
+        Returns:
+            Dict mapping feature key to error message
+        """
+        events_lazy = self.get_migration_events(migration_id)
+        events_df = (
+            events_lazy.filter(
+                (nw.col("event_type") == "feature_completed")
+                & (nw.col("error_message") != "")
+            )
+            .collect()
+            .to_polars()
+        )
+
+        result = {}
+        for row in events_df.iter_rows(named=True):
+            result[row["feature_key"]] = row["error_message"]
+
+        return result

metaxy/migrations/__init__.py

@@ -0,0 +1,34 @@
+"""Migration system for metadata version updates."""
+
+from metaxy.metadata_store.system_tables import SystemTableStorage
+from metaxy.migrations.detector import detect_migration
+from metaxy.migrations.executor import MigrationExecutor
+from metaxy.migrations.models import (
+    CustomMigration,
+    DiffMigration,
+    FullGraphMigration,
+    Migration,
+    MigrationResult,
+)
+from metaxy.migrations.ops import (
+    BaseOperation,
+    DataVersionReconciliation,
+    MetadataBackfill,
+)
+
+__all__ = [
+    # Core migration types
+    "Migration",
+    "DiffMigration",
+    "FullGraphMigration",
+    "CustomMigration",
+    "MigrationResult",
+    # Operations (for custom migrations)
+    "BaseOperation",
+    "DataVersionReconciliation",
+    "MetadataBackfill",
+    # Migration workflow
+    "detect_migration",
+    "MigrationExecutor",
+    "SystemTableStorage",
+]

metaxy/migrations/detector.py

@@ -0,0 +1,153 @@
+"""Feature change detection for automatic migration generation."""
+
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from metaxy.graph.diff.differ import GraphDiffer
+from metaxy.models.feature import FeatureGraph
+
+if TYPE_CHECKING:
+    from metaxy.metadata_store.base import MetadataStore
+    from metaxy.migrations.models import DiffMigration
+
+
+def detect_migration(
+    store: "MetadataStore",
+    from_snapshot_version: str | None = None,
+    ops: list[dict[str, Any]] | None = None,
+    migrations_dir: Path | None = None,
+    name: str | None = None,
+) -> "DiffMigration | None":
+    """Detect migration needed between snapshots and write YAML file.
+
+    Compares the latest snapshot in the store (or specified from_snapshot_version)
+    with the current active graph to detect changes and generate a migration YAML file.
+
+    Args:
+        store: Metadata store containing snapshot metadata
+        from_snapshot_version: Source snapshot version (defaults to latest in store)
+        ops: List of operation dicts with "type" field (defaults to [{"type": "metaxy.migrations.ops.DataVersionReconciliation"}])
+        migrations_dir: Directory to write migration YAML (defaults to .metaxy/migrations/)
+        name: Migration name (creates {timestamp}_{name} ID and filename)
+
+    Returns:
+        DiffMigration if changes detected and written, None otherwise
+
+    Example:
+        >>> # Compare latest snapshot in store vs current graph
+        >>> with store:
+        ...     migration = detect_migration(store)
+        ...     if migration:
+        ...         print(f"Migration written to {migration.yaml_path}")
+
+        >>> # Use custom operation
+        >>> migration = detect_migration(store, ops=[{"type": "myproject.ops.CustomOp"}])
+
+        >>> # Use custom name
+        >>> migration = detect_migration(store, name="example_migration")
+    """
+    from metaxy.migrations.models import DiffMigration
+
+    differ = GraphDiffer()
+
+    # Get from_snapshot_version (use latest if not specified)
+    if from_snapshot_version is None:
+        snapshots = store.read_graph_snapshots()
+        if snapshots.height == 0:
+            # No snapshots in store - nothing to migrate from
+            return None
+        from_snapshot_version = snapshots["snapshot_version"][0]
+
+    # At this point, from_snapshot_version is guaranteed to be a str
+    assert from_snapshot_version is not None  # Type narrowing for type checker
+
+    # Get to_snapshot_version from current active graph
+    active_graph = FeatureGraph.get_active()
+    if len(active_graph.features_by_key) == 0:
+        # No features in active graph - nothing to migrate to
+        return None
+
+    to_snapshot_version = active_graph.snapshot_version
+
+    # Check if versions are the same (no changes)
+    if from_snapshot_version == to_snapshot_version:
+        return None
+
+    # Load snapshot data using GraphDiffer
+    try:
+        from_snapshot_data = differ.load_snapshot_data(store, from_snapshot_version)
+    except ValueError:
+        # Snapshot not found - nothing to migrate from
+        return None
+
+    # Build snapshot data for to_snapshot (current graph)
+    to_snapshot_data = active_graph.to_snapshot()
+
+    # Compute GraphDiff using GraphDiffer
+    graph_diff = differ.diff(
+        from_snapshot_data,
+        to_snapshot_data,
+        from_snapshot_version,
+        to_snapshot_version,
+    )
+
+    # Check if there are any changes
+    if not graph_diff.has_changes:
+        return None
+
+    # Generate migration ID (timestamp first for sorting)
+    timestamp = datetime.now(timezone.utc)
+    timestamp_str = timestamp.strftime("%Y%m%d_%H%M%S")
+    if name is not None:
+        migration_id = f"{timestamp_str}_{name}"
+    else:
+        migration_id = f"{timestamp_str}"
+
+    # ops is required - caller must specify
+    if ops is None:
+        raise ValueError(
+            "ops parameter is required - must explicitly specify migration operations. "
+            "Example: ops=[{'type': 'metaxy.migrations.ops.DataVersionReconciliation'}]"
+        )
+
+    # Default migrations directory
+    if migrations_dir is None:
+        migrations_dir = Path(".metaxy/migrations")
+
+    migrations_dir.mkdir(parents=True, exist_ok=True)
+
+    # Find parent migration (latest migration in chain)
+    from metaxy.migrations.loader import find_latest_migration
+
+    parent = find_latest_migration(migrations_dir)
+    if parent is None:
+        parent = "initial"
+
+    # Create minimal DiffMigration - affected_features and description are computed on-demand
+    migration = DiffMigration(
+        migration_id=migration_id,
+        created_at=timestamp,
+        parent=parent,
+        from_snapshot_version=from_snapshot_version,
+        to_snapshot_version=to_snapshot_version,
+        ops=ops,
+    )
+
+    # Write migration YAML file
+    import yaml
+
+    yaml_path = migrations_dir / f"{migration_id}.yaml"
+    migration_yaml = {
+        "id": migration.migration_id,
+        "created_at": migration.created_at.isoformat(),
+        "parent": migration.parent,
+        "from_snapshot_version": migration.from_snapshot_version,
+        "to_snapshot_version": migration.to_snapshot_version,
+        "ops": migration.ops,
+    }
+
+    with open(yaml_path, "w") as f:
+        yaml.safe_dump(migration_yaml, f, sort_keys=False, default_flow_style=False)
+
+    return migration