metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/migrations/generator.py

@@ -0,0 +1,319 @@
+"""Migration generation."""
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+import narwhals as nw
+
+from metaxy.graph.diff.differ import GraphDiffer
+from metaxy.metadata_store.exceptions import FeatureNotFoundError
+from metaxy.metadata_store.system import SystemTableStorage
+from metaxy.migrations.models import DiffMigration
+from metaxy.migrations.ops import DataVersionReconciliation
+from metaxy.models.types import FeatureKey
+
+if TYPE_CHECKING:
+    from metaxy.metadata_store.base import MetadataStore
+    from metaxy.models.feature import FeatureGraph
+
+
+def _is_upstream_of(
+    upstream_key: FeatureKey, downstream_key: FeatureKey, graph: "FeatureGraph"
+) -> bool:
+    """Check if upstream_key is in the dependency chain of downstream_key.
+
+    Args:
+        upstream_key: Potential upstream feature
+        downstream_key: Feature to check dependencies for
+        graph: Feature graph
+
+    Returns:
+        True if upstream_key is a direct or transitive dependency of downstream_key
+    """
+    plan = graph.get_feature_plan(downstream_key)
+
+    if plan.deps is None:
+        return False
+
+    # Check direct dependencies
+    for dep in plan.deps:
+        if dep.key == upstream_key:
+            return True
+
+    # Check transitive dependencies (recursive)
+    for dep in plan.deps:
+        if _is_upstream_of(upstream_key, dep.key, graph):
+            return True
+
+    return False
+
+
+def generate_migration(
+    store: "MetadataStore",
+    *,
+    project: str,
+    from_snapshot_version: str | None = None,
+    to_snapshot_version: str | None = None,
+    class_path_overrides: dict[str, str] | None = None,
+) -> DiffMigration | None:
+    """Generate migration from detected feature changes or between snapshots.
+
+    Two modes of operation:
+
+    1. **Default mode** (both snapshot_versions None):
+       - Compares latest recorded snapshot (store) vs current active graph (code)
+       - This is the normal workflow: detect code changes
+
+    2. **Historical mode** (both snapshot_versions provided):
+       - Reconstructs from_graph from from_snapshot_version
+       - Reconstructs to_graph from to_snapshot_version
+       - Compares these two historical registries
+       - Useful for: backfilling migrations, testing, recovery
+
+    Generates explicit operations for ALL affected features (root + downstream).
+    Each downstream feature gets its own DataVersionReconciliation operation.
+
+    Args:
+        store: Metadata store to check
+        project: Project name for filtering snapshots
+        from_snapshot_version: Optional snapshot version to compare from (historical mode)
+        to_snapshot_version: Optional snapshot version to compare to (historical mode)
+        class_path_overrides: Optional overrides for moved/renamed feature classes
+
+    Returns:
+        Migration object, or None if no changes detected
+
+    Raises:
+        ValueError: If only one snapshot_version is provided, or snapshots not found
+
+    Example (default mode):
+        ```py
+        migration = generate_migration(store, project="my_project")
+        if migration:
+        migration.to_yaml("migrations/001_update.yaml")
+
+        ```
+    Example (historical mode):
+        ```py
+        migration = generate_migration(
+            store,
+            project="my_project",
+            from_snapshot_version="abc123...",
+            to_snapshot_version="def456...",
+            )
+        ```
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    if from_snapshot_version is None:
+        # Default mode: get from store's latest snapshot
+        from metaxy.metadata_store.system.keys import FEATURE_VERSIONS_KEY
+
+        try:
+            feature_versions = store.read_metadata(
+                FEATURE_VERSIONS_KEY, current_only=False
+            )
+            # Get most recent snapshot - only collect the top row
+            latest_snapshot = nw.from_native(
+                feature_versions.sort("recorded_at", descending=True).head(1).collect()
+            )
+            if latest_snapshot.shape[0] > 0:
+                from_snapshot_version = latest_snapshot["metaxy_snapshot_version"][0]
+                print(f"From: latest snapshot {from_snapshot_version}...")
+            else:
+                raise ValueError(
+                    "No feature graph snapshot found in metadata store. "
+                    "Run 'metaxy graph push' first to record feature versions before generating migrations."
+                )
+        except FeatureNotFoundError:
+            raise ValueError(
+                "No feature versions recorded yet. "
+                "Run 'metaxy graph push' first to record the feature graph snapshot."
+            )
+    else:
+        print(f"From: snapshot {from_snapshot_version}...")
+
+    # Step 2: Determine to_graph and to_snapshot_version
+    if to_snapshot_version is None:
+        # Default mode: record current active graph and use its snapshot
+        # This ensures the to_snapshot is available in the store for comparison
+        snapshot_result = SystemTableStorage(store).push_graph_snapshot()
+        to_snapshot_version = snapshot_result.snapshot_version
+        was_already_pushed = snapshot_result.already_pushed
+        to_graph = FeatureGraph.get_active()
+        if was_already_pushed:
+            print(
+                f"To: current active graph (snapshot {to_snapshot_version}... already pushed)"
+            )
+        else:
+            print(
+                f"To: current active graph (snapshot {to_snapshot_version}... pushed)"
+            )
+
+    else:
+        # Historical mode: load from snapshot with force_reload
+        # force_reload ensures we get current code from disk, not cached imports
+        to_graph = SystemTableStorage(store).load_graph_from_snapshot(
+            snapshot_version=to_snapshot_version,
+            class_path_overrides=class_path_overrides,
+            force_reload=True,
+        )
+        print(f"To: snapshot {to_snapshot_version}...")
+
+    # Step 3: Detect changes by comparing snapshot_versions directly
+    # We don't reconstruct from_graph - just compare snapshot_versions from the store
+    # This avoids issues with stale cached imports when files have changed
+    assert from_snapshot_version is not None, "from_snapshot_version must be set by now"
+    assert to_snapshot_version is not None, "to_snapshot_version must be set by now"
+
+    # Use GraphDiffer to detect changes
+    differ = GraphDiffer()
+
+    # 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
+        print("No from_snapshot found in store.")
+        return None
+
+    # Build snapshot data for to_snapshot
+    to_snapshot_data = to_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:
+        print("No feature changes detected. All features up to date!")
+        return None
+
+    # Create operations for root changed features
+    root_operations = []
+    for node in graph_diff.changed_nodes:
+        feature_key_str = node.feature_key.to_string()
+        feature_key_str.replace("/", "_")
+
+        root_operations.append(DataVersionReconciliation())
+
+    if not root_operations:
+        print("No feature changes detected. All features up to date!")
+        return None
+
+    # Generate migration ID and timestamp
+    timestamp = datetime.now()
+    timestamp_str = timestamp.strftime("%Y%m%d_%H%M%S")
+    migration_id = f"migration_{timestamp_str}"
+
+    # Show detected root changes
+    print(f"\nDetected {len(root_operations)} root feature change(s):")
+    for op in root_operations:
+        feature_key_str = FeatureKey(op.feature_key).to_string()
+        print(f"  ✓ {feature_key_str}")
+
+    # Discover downstream features that need reconciliation (use to_graph)
+    root_keys = [FeatureKey(op.feature_key) for op in root_operations]
+    downstream_keys = to_graph.get_downstream_features(root_keys)
+
+    # Create explicit operations for downstream features
+    downstream_operations = []
+
+    if downstream_keys:
+        print(
+            f"\nGenerating explicit operations for {len(downstream_keys)} downstream feature(s):"
+        )
+
+    for downstream_key in downstream_keys:
+        feature_key_str = downstream_key.to_string()
+        feature_cls = to_graph.features_by_key[downstream_key]
+
+        # Check if feature exists in from_snapshot (if not, it's new - skip)
+        try:
+            from_metadata = store.read_metadata(
+                feature_cls,
+                current_only=False,
+                allow_fallback=False,
+                filters=[nw.col("metaxy_snapshot_version") == from_snapshot_version],
+            )
+            # Only collect head(1) to check existence
+            from_metadata_sample = nw.from_native(from_metadata.head(1).collect())
+            if from_metadata_sample.shape[0] == 0:
+                # Feature doesn't exist in from_snapshot - it's new, skip
+                print(f"  ⊘ {feature_key_str} (new feature, skipping)")
+                continue
+        except FeatureNotFoundError:
+            # Feature not materialized yet
+            print(f"  ⊘ {feature_key_str} (not materialized yet, skipping)")
+            continue
+
+        # Determine which root changes affect this downstream feature
+        to_graph.get_feature_plan(downstream_key)
+        affected_by = []
+
+        for root_op in root_operations:
+            root_key = FeatureKey(root_op.feature_key)
+            # Check if this root is in the upstream dependency chain
+            if _is_upstream_of(root_key, downstream_key, to_graph):
+                affected_by.append(root_key.to_string())
+
+        # Build informative reason
+        if len(affected_by) == 1:
+            f"Reconcile field_provenance due to changes in: {affected_by[0]}"
+        else:
+            (f"Reconcile field_provenance due to changes in: {', '.join(affected_by)}")
+
+        # Create operation (feature versions derived from snapshots)
+        # DataVersionReconciliation doesn't have id, feature_key, or reason params
+        # It only has a type field since it applies to all affected features
+        downstream_operations.append(DataVersionReconciliation())
+
+        print(f"  ✓ {feature_key_str}")
+
+    # Combine all operations
+    all_operations = root_operations + downstream_operations
+
+    print(
+        f"\nGenerated {len(all_operations)} total operations "
+        f"({len(root_operations)} root + {len(downstream_operations)} downstream)"
+    )
+
+    # Find the latest migration to set as parent
+    from metaxy.metadata_store.system import EVENTS_KEY
+
+    parent_migration_id = None
+    try:
+        existing_migrations = store.read_metadata(EVENTS_KEY, current_only=False)
+        # Get most recent migration by timestamp - only collect the top row
+        latest = nw.from_native(
+            existing_migrations.sort("timestamp", descending=True).head(1).collect()
+        )
+        if latest.shape[0] > 0:
+            parent_migration_id = latest["migration_id"][0]
+    except FeatureNotFoundError:
+        # No migrations yet
+        pass
+
+    # Note: from_snapshot_version and to_snapshot_version were already resolved earlier
+
+    # Create migration (serialize operations to dicts)
+    len(root_operations)
+
+    # DiffMigration expects 'ops' as list of dicts with 'type' field
+    # Since all operations are DataVersionReconciliation, create a single operation dict
+    ops = [{"type": "metaxy.migrations.ops.DataVersionReconciliation"}]
+
+    migration = DiffMigration(
+        migration_id=migration_id,
+        parent=parent_migration_id or "initial",
+        from_snapshot_version=from_snapshot_version,
+        to_snapshot_version=to_snapshot_version,
+        created_at=timestamp,
+        ops=ops,
+    )
+
+    return migration

metaxy/migrations/loader.py

@@ -0,0 +1,231 @@
+"""Load migrations from YAML files."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from metaxy.migrations.models import Migration
+
+
+def load_migration_from_yaml(yaml_path: Path) -> "Migration":
+    """Load migration from YAML file.
+
+    Uses Pydantic's discriminated unions for automatic polymorphic deserialization
+    based on the migration_type field.
+
+    Args:
+        yaml_path: Path to migration YAML file
+
+    Returns:
+        Migration instance (DiffMigration or FullGraphMigration)
+
+    Raises:
+        FileNotFoundError: If YAML file doesn't exist
+        ValueError: If YAML is invalid or migration type is not supported
+    """
+    import yaml
+
+    from metaxy.migrations.models import MigrationAdapter
+
+    if not yaml_path.exists():
+        raise FileNotFoundError(f"Migration YAML not found: {yaml_path}")
+
+    with open(yaml_path) as f:
+        data = yaml.safe_load(f)
+
+    # Use Pydantic's discriminated union to automatically deserialize
+    try:
+        migration = MigrationAdapter.validate_python(data)
+    except Exception as e:
+        raise ValueError(f"Failed to load migration from {yaml_path}: {e}") from e
+
+    return migration
+
+
+def find_migration_yaml(migration_id: str, migrations_dir: Path | None = None) -> Path:
+    """Find YAML file for a migration ID by searching all YAML files.
+
+    Args:
+        migration_id: Migration ID (e.g., "20250127_120000" or "20250127_120000_feature_update")
+        migrations_dir: Directory containing migrations (defaults to .metaxy/migrations/)
+
+    Returns:
+        Path to migration YAML file
+
+    Raises:
+        FileNotFoundError: If migration YAML not found
+    """
+    if migrations_dir is None:
+        migrations_dir = Path(".metaxy/migrations")
+
+    if not migrations_dir.exists():
+        raise FileNotFoundError(
+            f"Migration '{migration_id}' not found. "
+            f"Migrations directory does not exist: {migrations_dir}"
+        )
+
+    # Search through all YAML files to find the one with matching ID
+    for yaml_file in migrations_dir.glob("*.yaml"):
+        try:
+            migration = load_migration_from_yaml(yaml_file)
+            if migration.migration_id == migration_id:
+                return yaml_file
+        except Exception:
+            # Skip files that can't be loaded
+            continue
+
+    # Not found - list available migrations
+    available = []
+    for yaml_file in migrations_dir.glob("*.yaml"):
+        try:
+            migration = load_migration_from_yaml(yaml_file)
+            available.append(migration.migration_id)
+        except Exception:
+            continue
+
+    raise FileNotFoundError(
+        f"Migration '{migration_id}' not found in {migrations_dir}.\n"
+        f"Available migrations: {available}"
+    )
+
+
+def list_migrations(migrations_dir: Path | None = None) -> list[str]:
+    """List all available migration IDs.
+
+    Args:
+        migrations_dir: Directory containing migrations (defaults to .metaxy/migrations/)
+
+    Returns:
+        List of migration IDs sorted by creation time
+    """
+    if migrations_dir is None:
+        migrations_dir = Path(".metaxy/migrations")
+
+    if not migrations_dir.exists():
+        return []
+
+    yaml_files = sorted(migrations_dir.glob("*.yaml"))
+    return [f.stem for f in yaml_files]
+
+
+def find_latest_migration(migrations_dir: Path | None = None) -> str | None:
+    """Find the latest migration ID (head of the chain).
+
+    Args:
+        migrations_dir: Directory containing migrations (defaults to .metaxy/migrations/)
+
+    Returns:
+        Migration ID of the head, or None if no migrations exist
+
+    Raises:
+        ValueError: If multiple heads detected (conflict)
+    """
+    from metaxy.migrations.models import Migration
+
+    if migrations_dir is None:
+        migrations_dir = Path(".metaxy/migrations")
+
+    if not migrations_dir.exists():
+        return None
+
+    # Load all migrations - all migrations form chains via parent IDs
+    migrations: dict[str, Migration] = {}
+    for yaml_file in migrations_dir.glob("*.yaml"):
+        migration = load_migration_from_yaml(yaml_file)
+        migrations[migration.migration_id] = migration
+
+    if not migrations:
+        return None
+
+    # Find migrations that are parents of others
+    all_parents = {m.parent for m in migrations.values() if m.parent != "initial"}
+
+    # Find heads (migrations that are not parents of any other migration)
+    heads = [mid for mid in migrations.keys() if mid not in all_parents]
+
+    if len(heads) == 0:
+        # This means there's a cycle or orphaned migrations
+        raise ValueError(
+            "No head migration found - possible cycle in migration chain. "
+            f"All migrations: {list(migrations.keys())}"
+        )
+
+    if len(heads) > 1:
+        raise ValueError(
+            f"Multiple migration heads detected: {heads}. "
+            "This usually means two migrations were created in parallel. "
+            "Please merge them by creating a new migration that depends on one head, "
+            "or delete one of the conflicting migrations."
+        )
+
+    return heads[0]
+
+
+def build_migration_chain(
+    migrations_dir: Path | None = None,
+) -> list["Migration"]:
+    """Build ordered migration chain from parent IDs.
+
+    Args:
+        migrations_dir: Directory containing migrations (defaults to .metaxy/migrations/)
+
+    Returns:
+        List of migrations in order from oldest to newest
+
+    Raises:
+        ValueError: If chain is invalid (cycles, orphans, multiple heads)
+    """
+    from metaxy.migrations.models import Migration
+
+    if migrations_dir is None:
+        migrations_dir = Path(".metaxy/migrations")
+
+    if not migrations_dir.exists():
+        return []
+
+    # Load all migrations - all migrations form chains via parent IDs
+    migrations: dict[str, Migration] = {}
+    for yaml_file in sorted(migrations_dir.glob("*.yaml")):
+        migration = load_migration_from_yaml(yaml_file)
+        migrations[migration.migration_id] = migration
+
+    if not migrations:
+        return []
+
+    # Validate single head
+    head_id = find_latest_migration(migrations_dir)
+    if head_id is None:
+        return []
+
+    # Build chain by following parent links backwards
+    chain = []
+    current_id: str | None = head_id
+
+    visited = set()
+    while current_id is not None and current_id != "initial":
+        if current_id in visited:
+            raise ValueError(f"Cycle detected in migration chain at: {current_id}")
+
+        if current_id not in migrations:
+            raise ValueError(
+                f"Migration '{current_id}' referenced as parent but YAML not found. "
+                f"Available migrations: {list(migrations.keys())}"
+            )
+
+        visited.add(current_id)
+        migration = migrations[current_id]
+        chain.append(migration)
+        current_id = migration.parent
+
+    # Reverse to get oldest-first order
+    chain.reverse()
+
+    # Validate all migrations are in the chain (no orphans)
+    if len(chain) != len(migrations):
+        orphans = set(migrations.keys()) - set(m.migration_id for m in chain)
+        raise ValueError(
+            f"Orphaned migrations detected (not in main chain): {orphans}. "
+            "Each migration must have parent pointing to previous migration or 'initial'."
+        )
+
+    return chain