dory-processor-sdk 0.0.1__py3-none-any.whl

dory/recovery/golden_snapshot.py

@@ -0,0 +1,713 @@
+"""
+Golden Snapshot Manager
+
+Captures and manages golden snapshots of processor state to prevent
+100% data loss during resets. Implements:
+- Snapshot capture with checksums
+- Versioned snapshot storage using StateManager
+- Snapshot validation
+- Restoration from snapshots
+
+Storage: Uses StateManager to persist snapshots to the same backend
+as processor state (ConfigMap in production, local file in development).
+This ensures snapshots survive pod restarts and node failures.
+"""
+
+import asyncio
+import hashlib
+import json
+import logging
+import time
+from dataclasses import dataclass, field, asdict
+from enum import Enum
+from typing import Any, Dict, Optional, List, Callable, TYPE_CHECKING
+import gzip
+import base64
+
+if TYPE_CHECKING:
+    from dory.migration.state_manager import StateManager
+
+logger = logging.getLogger(__name__)
+
+
+class SnapshotStorageError(Exception):
+    """Raised when snapshot storage operations fail."""
+    pass
+
+
+class SnapshotValidationError(Exception):
+    """Raised when snapshot validation fails."""
+    pass
+
+
+class SnapshotFormat(Enum):
+    """Snapshot storage format."""
+    JSON = "json"
+    JSON_GZ = "json.gz"  # Compressed JSON
+    BINARY = "binary"
+
+
+@dataclass
+class SnapshotMetadata:
+    """
+    Metadata about a golden snapshot.
+
+    Includes version, timestamps, checksums, and size information.
+    """
+    snapshot_id: str
+    processor_id: str
+    created_at: float
+    state_version: str
+    checksum: str
+    size_bytes: int
+    compressed: bool = False
+    format: str = "json"
+    validation_passed: bool = True
+    restore_count: int = 0
+    last_restored_at: Optional[float] = None
+    tags: Dict[str, str] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SnapshotMetadata":
+        """Create from dictionary."""
+        return cls(**data)
+
+    def age_seconds(self) -> float:
+        """Get age of snapshot in seconds."""
+        return time.time() - self.created_at
+
+
+@dataclass
+class Snapshot:
+    """
+    Complete snapshot including metadata and state data.
+    """
+    metadata: SnapshotMetadata
+    state_data: Dict[str, Any]
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "metadata": self.metadata.to_dict(),
+            "state_data": self.state_data
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Snapshot":
+        """Create from dictionary."""
+        return cls(
+            metadata=SnapshotMetadata.from_dict(data["metadata"]),
+            state_data=data["state_data"]
+        )
+
+
+class GoldenSnapshotManager:
+    """
+    Manages golden snapshots of processor state.
+
+    Features:
+    - Automatic snapshot capture at key points
+    - Checksum validation
+    - Multiple snapshot versions
+    - Compression support
+    - Restore with validation
+    - Snapshot lifecycle management
+    - Persistent storage via StateManager (ConfigMap/S3/PVC)
+
+    Usage:
+        from dory.migration.state_manager import StateManager
+
+        state_manager = StateManager(backend="configmap")
+        manager = GoldenSnapshotManager(state_manager=state_manager)
+
+        # Capture snapshot
+        snapshot = await manager.capture_snapshot(
+            processor_id="my-processor",
+            state_data={"key": "value"},
+            tags={"version": "1.0"}
+        )
+
+        # List snapshots
+        snapshots = await manager.list_snapshots(processor_id="my-processor")
+
+        # Restore from snapshot
+        state = await manager.restore_snapshot(snapshot.metadata.snapshot_id)
+    """
+
+    # Key prefix for snapshot storage
+    SNAPSHOT_KEY_PREFIX = "snapshot"
+
+    def __init__(
+        self,
+        state_manager: "StateManager",
+        max_snapshots_per_processor: int = 5,
+        compression_enabled: bool = True,
+        checksum_algorithm: str = "sha256",
+        auto_cleanup: bool = True,
+        on_capture: Optional[Callable] = None,
+        on_restore: Optional[Callable] = None,
+    ):
+        """
+        Initialize golden snapshot manager.
+
+        Args:
+            state_manager: StateManager instance for persistent storage
+            max_snapshots_per_processor: Maximum snapshots per processor
+            compression_enabled: Whether to compress snapshots
+            checksum_algorithm: Algorithm for checksums (sha256, md5)
+            auto_cleanup: Automatically cleanup old snapshots
+            on_capture: Callback when snapshot is captured
+            on_restore: Callback when snapshot is restored
+        """
+        self._state_manager = state_manager
+        self.max_snapshots_per_processor = max_snapshots_per_processor
+        self.compression_enabled = compression_enabled
+        self.checksum_algorithm = checksum_algorithm
+        self.auto_cleanup = auto_cleanup
+        self.on_capture = on_capture
+        self.on_restore = on_restore
+
+        # In-memory cache of snapshot metadata (loaded on first access)
+        self._snapshot_cache: Dict[str, List[Dict[str, Any]]] = {}
+        self._cache_lock = asyncio.Lock()
+
+        # Metrics
+        self._capture_count = 0
+        self._restore_count = 0
+        self._validation_failures = 0
+
+        logger.info(
+            f"GoldenSnapshotManager initialized: "
+            f"compression={compression_enabled}, max_per_processor={max_snapshots_per_processor}"
+        )
+
+    async def capture_snapshot(
+        self,
+        processor_id: str,
+        state_data: Dict[str, Any],
+        state_version: str = "1.0",
+        tags: Optional[Dict[str, str]] = None,
+        validate_before_save: bool = True,
+    ) -> Snapshot:
+        """
+        Capture a golden snapshot of processor state.
+
+        Args:
+            processor_id: ID of processor
+            state_data: State data to snapshot
+            state_version: Version of state schema
+            tags: Optional tags for the snapshot
+            validate_before_save: Validate state before saving
+
+        Returns:
+            Captured snapshot with metadata
+
+        Raises:
+            SnapshotValidationError: If validation fails
+            SnapshotStorageError: If storage fails
+        """
+        logger.info(f"Capturing snapshot for processor {processor_id}")
+
+        # Validate state data
+        if validate_before_save:
+            if not self._validate_state_data(state_data):
+                raise SnapshotValidationError("State data validation failed")
+
+        # Generate snapshot ID
+        snapshot_id = self._generate_snapshot_id(processor_id)
+
+        # Serialize state data
+        state_json = json.dumps(state_data, sort_keys=True)
+
+        # Compress if enabled
+        if self.compression_enabled:
+            state_bytes = gzip.compress(state_json.encode())
+            compressed = True
+            format_type = SnapshotFormat.JSON_GZ.value
+        else:
+            state_bytes = state_json.encode()
+            compressed = False
+            format_type = SnapshotFormat.JSON.value
+
+        # Calculate checksum
+        checksum = self._calculate_checksum(state_bytes)
+
+        # Create metadata
+        metadata = SnapshotMetadata(
+            snapshot_id=snapshot_id,
+            processor_id=processor_id,
+            created_at=time.time(),
+            state_version=state_version,
+            checksum=checksum,
+            size_bytes=len(state_bytes),
+            compressed=compressed,
+            format=format_type,
+            tags=tags or {},
+        )
+
+        # Create snapshot
+        snapshot = Snapshot(metadata=metadata, state_data=state_data)
+
+        # Save to storage
+        try:
+            await self._save_snapshot(snapshot, state_bytes)
+        except Exception as e:
+            logger.error(f"Failed to save snapshot: {e}")
+            raise SnapshotStorageError(f"Failed to save snapshot: {e}")
+
+        # Update metrics
+        self._capture_count += 1
+
+        # Cleanup old snapshots if auto_cleanup enabled
+        if self.auto_cleanup:
+            await self._cleanup_old_snapshots(processor_id)
+
+        # Call capture callback
+        if self.on_capture:
+            try:
+                if asyncio.iscoroutinefunction(self.on_capture):
+                    await self.on_capture(snapshot)
+                else:
+                    self.on_capture(snapshot)
+            except Exception as e:
+                logger.warning(f"Capture callback failed: {e}")
+
+        logger.info(
+            f"Snapshot captured: id={snapshot_id}, size={len(state_bytes)} bytes, "
+            f"compressed={compressed}, checksum={checksum[:8]}..."
+        )
+
+        return snapshot
+
+    async def restore_snapshot(
+        self,
+        snapshot_id: str,
+        validate_checksum: bool = True,
+        update_metadata: bool = True,
+    ) -> Dict[str, Any]:
+        """
+        Restore state from a snapshot.
+
+        Args:
+            snapshot_id: ID of snapshot to restore
+            validate_checksum: Validate checksum before restoring
+            update_metadata: Update metadata (restore count, timestamp)
+
+        Returns:
+            Restored state data
+
+        Raises:
+            SnapshotValidationError: If validation fails
+            SnapshotStorageError: If snapshot not found or load fails
+        """
+        logger.info(f"Restoring snapshot {snapshot_id}")
+
+        # Load snapshot
+        try:
+            snapshot, state_bytes = await self._load_snapshot(snapshot_id)
+        except Exception as e:
+            logger.error(f"Failed to load snapshot: {e}")
+            raise SnapshotStorageError(f"Failed to load snapshot: {e}")
+
+        # Validate checksum
+        if validate_checksum:
+            calculated_checksum = self._calculate_checksum(state_bytes)
+            if calculated_checksum != snapshot.metadata.checksum:
+                self._validation_failures += 1
+                raise SnapshotValidationError(
+                    f"Checksum mismatch: expected {snapshot.metadata.checksum}, "
+                    f"got {calculated_checksum}"
+                )
+
+        # Update metadata
+        if update_metadata:
+            snapshot.metadata.restore_count += 1
+            snapshot.metadata.last_restored_at = time.time()
+            await self._update_metadata(snapshot.metadata)
+
+        # Update metrics
+        self._restore_count += 1
+
+        # Call restore callback
+        if self.on_restore:
+            try:
+                if asyncio.iscoroutinefunction(self.on_restore):
+                    await self.on_restore(snapshot)
+                else:
+                    self.on_restore(snapshot)
+            except Exception as e:
+                logger.warning(f"Restore callback failed: {e}")
+
+        logger.info(
+            f"Snapshot restored: id={snapshot_id}, "
+            f"restore_count={snapshot.metadata.restore_count}"
+        )
+
+        return snapshot.state_data
+
+    async def list_snapshots(
+        self,
+        processor_id: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> List[SnapshotMetadata]:
+        """
+        List available snapshots.
+
+        Args:
+            processor_id: Filter by processor ID (required - StateManager requires processor ID)
+            limit: Maximum number of snapshots to return
+
+        Returns:
+            List of snapshot metadata, sorted by created_at (newest first)
+
+        Raises:
+            ValueError: If processor_id is not provided
+        """
+        if processor_id is None:
+            raise ValueError("processor_id is required for listing snapshots")
+
+        # Load snapshots from StateManager
+        snapshot_entries = await self._load_all_snapshots(processor_id)
+
+        # Extract metadata
+        snapshots = []
+        for entry in snapshot_entries:
+            try:
+                metadata = SnapshotMetadata.from_dict(entry["metadata"])
+                snapshots.append(metadata)
+            except Exception as e:
+                logger.warning(f"Failed to parse snapshot metadata: {e}")
+
+        # Already sorted by created_at (newest first) in _load_all_snapshots
+        # Apply limit
+        if limit is not None:
+            snapshots = snapshots[:limit]
+
+        return snapshots
+
+    async def delete_snapshot(self, snapshot_id: str) -> bool:
+        """
+        Delete a snapshot.
+
+        Args:
+            snapshot_id: ID of snapshot to delete
+
+        Returns:
+            True if deleted successfully
+        """
+        logger.info(f"Deleting snapshot {snapshot_id}")
+
+        try:
+            # Extract processor_id from snapshot_id (format: {processor_id}_{timestamp})
+            parts = snapshot_id.rsplit("_", 1)
+            if len(parts) != 2:
+                logger.error(f"Invalid snapshot ID format: {snapshot_id}")
+                return False
+
+            processor_id = parts[0]
+
+            # Load all snapshots
+            snapshots = await self._load_all_snapshots(processor_id)
+
+            # Find and remove the snapshot
+            original_count = len(snapshots)
+            snapshots = [
+                s for s in snapshots
+                if s["metadata"]["snapshot_id"] != snapshot_id
+            ]
+
+            if len(snapshots) == original_count:
+                logger.warning(f"Snapshot not found: {snapshot_id}")
+                return False
+
+            # Save back (or delete if no snapshots left)
+            if snapshots:
+                await self._save_all_snapshots(processor_id, snapshots)
+            else:
+                # No snapshots left, delete the entire state entry
+                storage_key = self._get_storage_key(processor_id)
+                await self._state_manager.delete_state(storage_key)
+                # Clear cache
+                self._snapshot_cache.pop(processor_id, None)
+
+            logger.info(f"Snapshot deleted: {snapshot_id}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to delete snapshot {snapshot_id}: {e}")
+            return False
+
+    async def get_latest_snapshot(
+        self,
+        processor_id: str
+    ) -> Optional[SnapshotMetadata]:
+        """
+        Get the latest snapshot for a processor.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            Latest snapshot metadata, or None if no snapshots exist
+        """
+        snapshots = await self.list_snapshots(processor_id=processor_id, limit=1)
+        return snapshots[0] if snapshots else None
+
+    async def delete_all_snapshots(self, processor_id: str) -> int:
+        """
+        Delete all snapshots for a processor.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            Number of snapshots deleted
+        """
+        logger.info(f"Deleting all snapshots for processor {processor_id}")
+
+        # Get count before deletion
+        snapshots = await self._load_all_snapshots(processor_id)
+        count = len(snapshots)
+
+        if count == 0:
+            return 0
+
+        # Delete the entire state entry
+        storage_key = self._get_storage_key(processor_id)
+        await self._state_manager.delete_state(storage_key)
+
+        # Clear cache
+        self._snapshot_cache.pop(processor_id, None)
+
+        logger.info(f"Deleted {count} snapshots for processor {processor_id}")
+        return count
+
+    def clear_cache(self, processor_id: Optional[str] = None) -> None:
+        """
+        Clear the in-memory snapshot cache.
+
+        Args:
+            processor_id: Clear cache for specific processor, or all if None
+        """
+        if processor_id:
+            self._snapshot_cache.pop(processor_id, None)
+        else:
+            self._snapshot_cache.clear()
+
+    def get_stats(self) -> Dict[str, Any]:
+        """
+        Get snapshot manager statistics.
+
+        Returns:
+            Dictionary of statistics
+        """
+        return {
+            "storage_backend": str(self._state_manager._backend),
+            "capture_count": self._capture_count,
+            "restore_count": self._restore_count,
+            "validation_failures": self._validation_failures,
+            "compression_enabled": self.compression_enabled,
+            "max_snapshots_per_processor": self.max_snapshots_per_processor,
+        }
+
+    # Private methods
+
+    def _get_storage_key(self, processor_id: str) -> str:
+        """Get the StateManager key for storing snapshots."""
+        return f"{processor_id}-{self.SNAPSHOT_KEY_PREFIX}"
+
+    def _generate_snapshot_id(self, processor_id: str) -> str:
+        """Generate unique snapshot ID."""
+        timestamp = int(time.time() * 1000)
+        return f"{processor_id}_{timestamp}"
+
+    def _calculate_checksum(self, data: bytes) -> str:
+        """Calculate checksum for data."""
+        if self.checksum_algorithm == "sha256":
+            return hashlib.sha256(data).hexdigest()
+        elif self.checksum_algorithm == "md5":
+            return hashlib.md5(data).hexdigest()
+        else:
+            raise ValueError(f"Unsupported checksum algorithm: {self.checksum_algorithm}")
+
+    def _validate_state_data(self, state_data: Dict[str, Any]) -> bool:
+        """
+        Validate state data before capture.
+
+        Args:
+            state_data: State data to validate
+
+        Returns:
+            True if valid
+        """
+        # Basic validation
+        if not isinstance(state_data, dict):
+            logger.error("State data must be a dictionary")
+            return False
+
+        # Check if serializable
+        try:
+            json.dumps(state_data)
+        except (TypeError, ValueError) as e:
+            logger.error(f"State data is not JSON serializable: {e}")
+            return False
+
+        return True
+
+    async def _load_all_snapshots(self, processor_id: str) -> List[Dict[str, Any]]:
+        """
+        Load all snapshots for a processor from StateManager.
+
+        Returns:
+            List of snapshot entries (each with 'metadata' and 'data_b64' keys)
+        """
+        async with self._cache_lock:
+            # Check cache first
+            if processor_id in self._snapshot_cache:
+                return self._snapshot_cache[processor_id]
+
+            storage_key = self._get_storage_key(processor_id)
+            state = await self._state_manager.load_state(storage_key)
+
+            if state is None:
+                self._snapshot_cache[processor_id] = []
+                return []
+
+            snapshots = state.get("snapshots", [])
+            self._snapshot_cache[processor_id] = snapshots
+            return snapshots
+
+    async def _save_all_snapshots(
+        self, processor_id: str, snapshots: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Save all snapshots for a processor to StateManager.
+
+        Args:
+            processor_id: Processor ID
+            snapshots: List of snapshot entries
+        """
+        storage_key = self._get_storage_key(processor_id)
+
+        state = {
+            "processor_id": processor_id,
+            "snapshot_count": len(snapshots),
+            "snapshots": snapshots,
+        }
+
+        await self._state_manager.save_state(storage_key, state)
+
+        # Update cache
+        async with self._cache_lock:
+            self._snapshot_cache[processor_id] = snapshots
+
+    async def _save_snapshot(self, snapshot: Snapshot, state_bytes: bytes) -> None:
+        """Save snapshot to StateManager."""
+        processor_id = snapshot.metadata.processor_id
+
+        # Load existing snapshots
+        snapshots = await self._load_all_snapshots(processor_id)
+
+        # Create snapshot entry with base64-encoded data
+        snapshot_entry = {
+            "metadata": snapshot.metadata.to_dict(),
+            "data_b64": base64.b64encode(state_bytes).decode("ascii"),
+        }
+
+        # Add new snapshot at the beginning (newest first)
+        snapshots.insert(0, snapshot_entry)
+
+        # Save back to StateManager
+        await self._save_all_snapshots(processor_id, snapshots)
+
+    async def _load_snapshot(self, snapshot_id: str) -> tuple[Snapshot, bytes]:
+        """Load snapshot from StateManager."""
+        # Extract processor_id from snapshot_id (format: {processor_id}_{timestamp})
+        parts = snapshot_id.rsplit("_", 1)
+        if len(parts) != 2:
+            raise SnapshotStorageError(f"Invalid snapshot ID format: {snapshot_id}")
+
+        processor_id = parts[0]
+
+        # Load all snapshots for this processor
+        snapshots = await self._load_all_snapshots(processor_id)
+
+        # Find the specific snapshot
+        snapshot_entry = None
+        for entry in snapshots:
+            if entry["metadata"]["snapshot_id"] == snapshot_id:
+                snapshot_entry = entry
+                break
+
+        if snapshot_entry is None:
+            raise SnapshotStorageError(f"Snapshot not found: {snapshot_id}")
+
+        # Decode data
+        state_bytes = base64.b64decode(snapshot_entry["data_b64"])
+        metadata = SnapshotMetadata.from_dict(snapshot_entry["metadata"])
+
+        # Decompress if needed
+        if metadata.compressed:
+            state_json = gzip.decompress(state_bytes).decode()
+        else:
+            state_json = state_bytes.decode()
+
+        # Parse JSON
+        state_data = json.loads(state_json)
+
+        # Create snapshot object
+        snapshot = Snapshot(metadata=metadata, state_data=state_data)
+
+        return snapshot, state_bytes
+
+    async def _update_metadata(self, metadata: SnapshotMetadata) -> None:
+        """Update snapshot metadata in StateManager."""
+        processor_id = metadata.processor_id
+
+        # Load all snapshots
+        snapshots = await self._load_all_snapshots(processor_id)
+
+        # Find and update the specific snapshot
+        for entry in snapshots:
+            if entry["metadata"]["snapshot_id"] == metadata.snapshot_id:
+                entry["metadata"] = metadata.to_dict()
+                break
+
+        # Save back
+        await self._save_all_snapshots(processor_id, snapshots)
+
+    async def _cleanup_old_snapshots(self, processor_id: str) -> int:
+        """
+        Cleanup old snapshots exceeding max limit.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            Number of snapshots deleted
+        """
+        snapshots = await self._load_all_snapshots(processor_id)
+
+        # Keep only max_snapshots_per_processor newest snapshots
+        if len(snapshots) <= self.max_snapshots_per_processor:
+            return 0
+
+        # Count snapshots to delete
+        deleted_count = len(snapshots) - self.max_snapshots_per_processor
+
+        # Keep only the newest snapshots
+        snapshots = snapshots[:self.max_snapshots_per_processor]
+
+        # Save back
+        await self._save_all_snapshots(processor_id, snapshots)
+
+        if deleted_count > 0:
+            logger.info(
+                f"Cleaned up {deleted_count} old snapshots for processor {processor_id}"
+            )
+
+        return deleted_count