metaxy 0.0.0__py3-none-any.whl

metaxy/migrations/models.py

@@ -0,0 +1,718 @@
+"""Type-safe migration models with Python class paths.
+
+Refactored migration system using:
+- Python class paths for polymorphic deserialization
+- Struct-based storage for graph data
+- Event-based status tracking
+"""
+
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+import pydantic
+from pydantic.types import AwareDatetime
+
+if TYPE_CHECKING:
+    from metaxy.graph.diff.diff_models import GraphDiff
+    from metaxy.metadata_store.base import MetadataStore
+
+
+class Migration(pydantic.BaseModel, ABC):  # pyright: ignore[reportUnsafeMultipleInheritance]
+    """Abstract base class for all migrations.
+
+    Subclasses must define:
+    - migration_type: Class path as Literal for polymorphic deserialization
+    - execute(): Migration logic
+
+    The migration_type field is used for storage and deserialization.
+    """
+
+    migration_id: str
+    created_at: AwareDatetime
+
+    @property
+    @abstractmethod
+    def migration_type(self) -> str:
+        """Get migration type (Python class path).
+
+        Returns:
+            Full Python class path (e.g., "metaxy.migrations.models.DiffMigration")
+        """
+        pass
+
+    @abstractmethod
+    def execute(
+        self,
+        store: "MetadataStore",
+        *,
+        dry_run: bool = False,
+    ) -> "MigrationResult":
+        """Execute the migration.
+
+        Args:
+            store: Metadata store to operate on
+            dry_run: If True, only validate without executing
+
+        Returns:
+            MigrationResult with execution details
+
+        Raises:
+            Exception: If migration fails
+        """
+        pass
+
+    @abstractmethod
+    def get_affected_features(self, store: "MetadataStore") -> list[str]:
+        """Get list of affected feature keys in topological order.
+
+        Args:
+            store: Metadata store for computing affected features
+
+        Returns:
+            List of feature key strings
+        """
+        pass
+
+    def to_storage_dict(self) -> dict[str, Any]:
+        """Convert to dict for storage.
+
+        Returns:
+            Dict with all fields including migration_type
+        """
+        data = self.model_dump(mode="python")
+        data["migration_type"] = self.migration_type
+        return data
+
+    @staticmethod
+    def from_storage_dict(data: dict[str, Any]) -> "Migration":
+        """Deserialize migration from storage dict.
+
+        Args:
+            data: Dict with migration_type and other fields
+
+        Returns:
+            Migration instance of appropriate subclass
+
+        Raises:
+            ValueError: If migration_type is invalid or class not found
+        """
+        migration_type = data.get("migration_type")
+        if not migration_type:
+            raise ValueError("Missing migration_type field")
+
+        # Dynamically import the class
+        try:
+            module_path, class_name = migration_type.rsplit(".", 1)
+            module = __import__(module_path, fromlist=[class_name])
+            cls = getattr(module, class_name)
+
+            if not issubclass(cls, Migration):
+                raise TypeError(
+                    f"{migration_type} must be a subclass of Migration, got {cls}"
+                )
+
+            return cls.model_validate(data)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to load migration class {migration_type}: {e}"
+            ) from e
+
+
+class DiffMigration(Migration):
+    """Migration based on graph diff between two snapshots.
+
+    Migrations form a chain via parent IDs (like git commits):
+    - id: Unique identifier for this migration
+    - parent: ID of parent migration ("initial" for first migration)
+    - from_snapshot_version: Source snapshot version
+    - to_snapshot_version: Target snapshot version
+    - ops: List of operation dicts with "type" field
+
+    The parent chain ensures migrations are applied in correct order.
+    Multiple heads (two migrations with no children) is an error.
+
+    All other information is computed on-demand:
+    - affected_features: Computed from GraphDiff when accessed
+    - operations: Instantiated from ops
+    - description: Auto-generated from affected features count
+
+    The graph diff is computed on-demand when needed using GraphDiffer.
+
+    Examples:
+        First migration:
+            DiffMigration(
+                migration_id="20250113_120000",
+                parent="initial",
+                from_snapshot_version="abc123...",
+                to_snapshot_version="def456...",
+                created_at=datetime.now(timezone.utc),
+            )
+
+        Subsequent migration:
+            DiffMigration(
+                migration_id="20250113_130000",
+                parent="20250113_120000",
+                from_snapshot_version="def456...",
+                to_snapshot_version="ghi789...",
+                created_at=datetime.now(timezone.utc),
+            )
+    """
+
+    # Stored fields - persisted to YAML in git
+    parent: str  # Parent migration ID or "initial"
+    from_snapshot_version: str
+    to_snapshot_version: str
+    ops: list[dict[str, Any]]  # Required - must explicitly specify operations
+
+    # Private attribute for caching computed graph diff
+    _graph_diff_cache: "GraphDiff | None" = pydantic.PrivateAttr(default=None)
+
+    @pydantic.model_validator(mode="before")
+    @classmethod
+    def deserialize_json_fields(cls, data: dict[str, Any]) -> dict[str, Any]:
+        """Deserialize JSON strings for ops (from storage).
+
+        Args:
+            data: Raw migration data
+
+        Returns:
+            Data with deserialized JSON fields
+        """
+        import json
+
+        data = dict(data)
+
+        # Deserialize ops from JSON string (from storage)
+        if isinstance(data.get("ops"), str):
+            data["ops"] = json.loads(data["ops"])
+
+        return data
+
+    @property
+    def migration_type(self) -> str:
+        """Get migration type."""
+        return "metaxy.migrations.models.DiffMigration"
+
+    def _get_graph_diff(self, store: "MetadataStore") -> "GraphDiff":
+        """Get or compute graph diff (cached).
+
+        Args:
+            store: Metadata store containing snapshots
+
+        Returns:
+            GraphDiff between snapshots
+        """
+        if self._graph_diff_cache is None:
+            self._graph_diff_cache = self.compute_graph_diff(store)
+        return self._graph_diff_cache
+
+    @property
+    def operations(self) -> list[Any]:
+        """Get operations for this migration.
+
+        Instantiates operations from stored ops (list of dicts with "type" field).
+
+        Returns:
+            List of operation instances
+        """
+        operations = []
+        for op_dict in self.ops:
+            op_type = op_dict.get("type")
+            if not op_type:
+                raise ValueError(f"Operation dict missing 'type' field: {op_dict}")
+            try:
+                # Dynamically import and instantiate the operation class
+                module_path, class_name = op_type.rsplit(".", 1)
+                module = __import__(module_path, fromlist=[class_name])
+                op_cls = getattr(module, class_name)
+                operations.append(op_cls())
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to instantiate operation {op_type}: {e}"
+                ) from e
+
+        return operations
+
+    @property
+    def description(self) -> str:
+        """Get auto-generated description for migration.
+
+        Returns:
+            Human-readable description based on affected features count
+        """
+        # Note: This accesses affected_features property which needs store access
+        # For display purposes, this is called after affected_features is computed
+        return self.auto_description
+
+    @property
+    def auto_description(self) -> str:
+        """Generate automatic description (requires store context).
+
+        Returns:
+            Human-readable description based on affected features
+        """
+        # This is used internally - callers should use get_description(store)
+        return "Migration: snapshot reconciliation"
+
+    def get_description(self, store: "MetadataStore") -> str:
+        """Get description for migration.
+
+        Args:
+            store: Metadata store for computing affected features
+
+        Returns:
+            Description string
+        """
+        affected = self.get_affected_features(store)
+        num_features = len(affected)
+        if num_features == 0:
+            return "No features affected"
+        elif num_features == 1:
+            return f"Migration: {affected[0]}"
+        else:
+            return f"Migration: {num_features} features affected"
+
+    def get_affected_features(self, store: "MetadataStore") -> list[str]:
+        """Get affected features in topological order (computed on-demand).
+
+        Args:
+            store: Metadata store containing snapshots (required for computation)
+
+        Returns:
+            List of feature key strings in topological order
+        """
+
+        graph_diff = self._get_graph_diff(store)
+
+        # Get changed feature keys (root changes)
+        changed_keys = {
+            node.feature_key.to_string() for node in graph_diff.changed_nodes
+        }
+
+        # Also include added nodes (though they typically don't have existing data to migrate)
+        for node in graph_diff.added_nodes:
+            changed_keys.add(node.feature_key.to_string())
+
+        # Build dependency map from the GraphDiff added/changed nodes
+        # We need to compute downstream dependencies to find all affected features
+        from metaxy.graph.diff.models import GraphData, GraphNode
+        from metaxy.graph.diff.traversal import GraphWalker
+        from metaxy.models.feature import FeatureGraph
+
+        # Get the active graph to extract dependencies
+        active_graph = FeatureGraph.get_active()
+
+        # Build GraphData from active graph for dependency analysis
+        nodes_dict = {}
+        for feature_key, feature_cls in active_graph.features_by_key.items():
+            plan = active_graph.get_feature_plan(feature_key)
+
+            # Extract dependencies from plan
+            dependencies = []
+            if plan.deps:
+                for dep in plan.deps:
+                    dependencies.append(dep.key)
+
+            nodes_dict[feature_key.to_string()] = GraphNode(
+                key=feature_key,
+                version=feature_cls.feature_version(),
+                dependencies=dependencies,
+            )
+
+        to_graph_data = GraphData(
+            nodes=nodes_dict, snapshot_version=self.to_snapshot_version
+        )
+
+        # Build reverse dependency map (feature -> dependents)
+        dependents_map: dict[str, set[str]] = {}
+        for node in to_graph_data.nodes.values():
+            for dep_key in node.dependencies:
+                dep_key_str = dep_key.to_string()
+                if dep_key_str not in dependents_map:
+                    dependents_map[dep_key_str] = set()
+                dependents_map[dep_key_str].add(node.key.to_string())
+
+        # Find all features affected (changed + their downstream)
+        affected = set(changed_keys)
+        queue = list(changed_keys)
+        while queue:
+            key_str = queue.pop(0)
+            if key_str in dependents_map:
+                for dependent in dependents_map[key_str]:
+                    if dependent not in affected:
+                        affected.add(dependent)
+                        queue.append(dependent)
+
+        # Get topological order for affected features
+        walker = GraphWalker(to_graph_data)
+        sorted_nodes = walker.topological_sort(nodes_to_include=affected)
+
+        return [node.key.to_string() for node in sorted_nodes]
+
+    def compute_graph_diff(self, store: "MetadataStore") -> "GraphDiff":
+        """Compute GraphDiff on-demand from snapshot versions.
+
+        Args:
+            store: Metadata store containing snapshots
+
+        Returns:
+            GraphDiff between from_snapshot_version and to_snapshot_version
+
+        Raises:
+            ValueError: If snapshots cannot be loaded
+        """
+        from metaxy.graph.diff.differ import GraphDiffer
+        from metaxy.models.feature import FeatureGraph
+
+        differ = GraphDiffer()
+
+        # Load from_snapshot data from store
+        from_snapshot_data = differ.load_snapshot_data(
+            store, self.from_snapshot_version
+        )
+
+        # Try to load to_snapshot from store, if it doesn't exist use active graph
+        try:
+            to_snapshot_data = differ.load_snapshot_data(
+                store, self.to_snapshot_version
+            )
+        except ValueError:
+            # Snapshot not recorded yet, use active graph
+            active_graph = FeatureGraph.get_active()
+            if active_graph.snapshot_version != self.to_snapshot_version:
+                raise ValueError(
+                    f"to_snapshot {self.to_snapshot_version} not found in store "
+                    f"and doesn't match active graph ({active_graph.snapshot_version})"
+                )
+            to_snapshot_data = active_graph.to_snapshot()
+
+        # Compute diff
+        return differ.diff(
+            from_snapshot_data,
+            to_snapshot_data,
+            self.from_snapshot_version,
+            self.to_snapshot_version,
+        )
+
+    def execute(
+        self,
+        store: "MetadataStore",
+        *,
+        dry_run: bool = False,
+    ) -> "MigrationResult":
+        """Execute diff-based migration.
+
+        Process:
+        1. Execute each operation in the operations list
+        2. For each operation:
+           - Check if feature already completed (resume support)
+           - Execute operation
+           - Record event
+        3. Return result
+
+        Args:
+            store: Metadata store
+            dry_run: If True, only validate
+
+        Returns:
+            MigrationResult
+        """
+        from metaxy.metadata_store.system_tables import SystemTableStorage
+
+        storage = SystemTableStorage(store)
+        start_time = datetime.now(timezone.utc)
+
+        if not dry_run:
+            # Write started event
+            storage.write_event(self.migration_id, "started")
+
+        affected_features_list = []
+        errors = {}
+        rows_affected_total = 0
+
+        # Execute operations (currently only DataVersionReconciliation is supported)
+        from metaxy.migrations.ops import DataVersionReconciliation
+
+        # Get affected features (computed on-demand)
+        affected_features_to_process = self.get_affected_features(store)
+
+        if len(self.operations) == 1 and isinstance(
+            self.operations[0], DataVersionReconciliation
+        ):
+            # DataVersionReconciliation applies to all affected features
+            op = self.operations[0]
+
+            for feature_key_str in affected_features_to_process:
+                # Check if already completed (resume support)
+                if not dry_run and storage.is_feature_completed(
+                    self.migration_id, feature_key_str
+                ):
+                    affected_features_list.append(feature_key_str)
+                    continue
+
+                # Log feature started
+                if not dry_run:
+                    storage.write_event(
+                        self.migration_id,
+                        "feature_started",
+                        feature_key=feature_key_str,
+                    )
+
+                try:
+                    # Execute operation for this feature
+                    rows_affected = op.execute_for_feature(
+                        store,
+                        feature_key_str,
+                        from_snapshot_version=self.from_snapshot_version,
+                        to_snapshot_version=self.to_snapshot_version,
+                        dry_run=dry_run,
+                    )
+
+                    # Log feature completed
+                    if not dry_run:
+                        storage.write_event(
+                            self.migration_id,
+                            "feature_completed",
+                            feature_key=feature_key_str,
+                            rows_affected=rows_affected,
+                        )
+
+                    affected_features_list.append(feature_key_str)
+                    rows_affected_total += rows_affected
+
+                except Exception as e:
+                    error_msg = str(e)
+                    errors[feature_key_str] = error_msg
+
+                    # Log feature failed
+                    if not dry_run:
+                        storage.write_event(
+                            self.migration_id,
+                            "feature_completed",
+                            feature_key=feature_key_str,
+                            error_message=error_msg,
+                        )
+
+                    continue
+        else:
+            # Future: Support other operation types here
+            raise NotImplementedError(
+                "Only DataVersionReconciliation is currently supported"
+            )
+
+        # Determine status
+        if dry_run:
+            status = "skipped"
+        elif len(errors) == 0:
+            status = "completed"
+            if not dry_run:
+                storage.write_event(self.migration_id, "completed")
+        else:
+            status = "failed"
+            if not dry_run:
+                storage.write_event(self.migration_id, "failed")
+
+        duration = (datetime.now(timezone.utc) - start_time).total_seconds()
+
+        return MigrationResult(
+            migration_id=self.migration_id,
+            status=status,
+            features_completed=len(affected_features_list),
+            features_failed=len(errors),
+            affected_features=affected_features_list,
+            errors=errors,
+            rows_affected=rows_affected_total,
+            duration_seconds=duration,
+            timestamp=start_time,
+        )
+
+
+class FullGraphMigration(Migration):
+    """Migration that operates within a single snapshot.
+
+    Used for operations that don't involve graph structure changes,
+    such as backfills or custom transformations on existing features.
+    """
+
+    snapshot_version: str
+    affected_features: list[str] = pydantic.Field(
+        default_factory=list
+    )  # Features to process
+    operations: list[Any] = pydantic.Field(default_factory=list)  # Custom operations
+    description: str | None = None
+    metadata: dict[str, Any] = pydantic.Field(default_factory=dict)
+
+    @pydantic.model_validator(mode="before")
+    @classmethod
+    def deserialize_json_fields(cls, data: dict[str, Any]) -> dict[str, Any]:
+        """Deserialize JSON strings for operations and metadata (from storage).
+
+        Args:
+            data: Raw migration data
+
+        Returns:
+            Data with deserialized JSON fields
+        """
+        import json
+
+        data = dict(data)
+
+        # Deserialize JSON strings (from storage)
+        if isinstance(data.get("operations"), str):
+            data["operations"] = json.loads(data["operations"])
+
+        if isinstance(data.get("metadata"), str):
+            data["metadata"] = json.loads(data["metadata"])
+
+        return data
+
+    @property
+    def migration_type(self) -> str:
+        """Get migration type."""
+        return "metaxy.migrations.models.FullGraphMigration"
+
+    def get_affected_features(self, store: "MetadataStore") -> list[str]:
+        """Get affected features.
+
+        Args:
+            store: Metadata store (not used for FullGraphMigration)
+
+        Returns:
+            List of feature key strings
+        """
+        return self.affected_features
+
+    def execute(
+        self,
+        store: "MetadataStore",
+        *,
+        dry_run: bool = False,
+    ) -> "MigrationResult":
+        """Execute full graph migration.
+
+        Subclasses should implement custom logic here.
+
+        Args:
+            store: Metadata store
+            dry_run: If True, only validate
+
+        Returns:
+            MigrationResult
+        """
+        # Base implementation: no-op
+        return MigrationResult(
+            migration_id=self.migration_id,
+            status="completed",
+            features_completed=0,
+            features_failed=0,
+            affected_features=[],
+            errors={},
+            rows_affected=0,
+            duration_seconds=0.0,
+            timestamp=datetime.now(timezone.utc),
+        )
+
+
+class CustomMigration(Migration):
+    """Base class for user-defined custom migrations.
+
+    Users can subclass this to implement completely custom migration logic.
+
+    Example:
+        class S3BackfillMigration(CustomMigration):
+            s3_bucket: str
+            s3_prefix: str
+
+            @property
+            def migration_type(self) -> str:
+                return "myproject.migrations.S3BackfillMigration"
+
+            def execute(self, store, *, dry_run=False):
+                # Custom logic here
+                ...
+    """
+
+    @property
+    def migration_type(self) -> str:
+        """Get migration type.
+
+        Subclasses must override this to return their full class path.
+        """
+        return f"{self.__class__.__module__}.{self.__class__.__name__}"
+
+    def get_affected_features(self, store: "MetadataStore") -> list[str]:
+        """Get affected features.
+
+        Args:
+            store: Metadata store (not used for CustomMigration base class)
+
+        Returns:
+            Empty list (subclasses should override)
+        """
+        return []
+
+    def execute(
+        self,
+        store: "MetadataStore",
+        *,
+        dry_run: bool = False,
+    ) -> "MigrationResult":
+        """Execute custom migration.
+
+        Subclasses must override this to implement custom logic.
+
+        Args:
+            store: Metadata store
+            dry_run: If True, only validate
+
+        Returns:
+            MigrationResult
+
+        Raises:
+            NotImplementedError: If not overridden by subclass
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} must implement execute() method"
+        )
+
+
+class MigrationResult(pydantic.BaseModel):
+    """Result of executing a migration."""
+
+    migration_id: str
+    status: str  # "completed", "failed", "skipped"
+    features_completed: int
+    features_failed: int
+    affected_features: list[str]
+    errors: dict[str, str]  # feature_key -> error message
+    rows_affected: int
+    duration_seconds: float
+    timestamp: AwareDatetime
+
+    def summary(self) -> str:
+        """Human-readable summary of migration result.
+
+        Returns:
+            Multi-line summary string
+        """
+        lines = [
+            f"Migration: {self.migration_id}",
+            f"Status: {self.status.upper()}",
+            f"Timestamp: {self.timestamp.isoformat()}",
+            f"Duration: {self.duration_seconds:.2f}s",
+            f"Features: {self.features_completed} completed, {self.features_failed} failed",
+            f"Rows affected: {self.rows_affected}",
+        ]
+
+        if self.affected_features:
+            lines.append("\nFeatures processed:")
+            for feature in self.affected_features:
+                lines.append(f"  ✓ {feature}")
+
+        if self.errors:
+            lines.append("\nErrors:")
+            for feature, error in self.errors.items():
+                lines.append(f"  ✗ {feature}: {error}")
+
+        return "\n".join(lines)