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

dory/migration/serialization.py

@@ -0,0 +1,135 @@
+"""
+State serialization utilities.
+
+Handles JSON serialization/deserialization with checksum validation.
+"""
+
+import hashlib
+import json
+import time
+from dataclasses import dataclass
+from typing import Any
+
+from dory.utils.errors import DoryStateError
+
+
+@dataclass
+class StateEnvelope:
+    """
+    Envelope wrapping state data with metadata.
+
+    Attributes:
+        payload: The actual state data
+        metadata: Metadata about when/where state was created
+        checksum: SHA256 checksum of payload for integrity
+    """
+
+    payload: dict[str, Any]
+    metadata: dict[str, Any]
+    checksum: str
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "payload": self.payload,
+            "metadata": self.metadata,
+            "checksum": self.checksum,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "StateEnvelope":
+        """Create from dictionary."""
+        return cls(
+            payload=data["payload"],
+            metadata=data["metadata"],
+            checksum=data["checksum"],
+        )
+
+
+class StateSerializer:
+    """
+    Serializes and deserializes state with integrity checking.
+
+    Uses JSON format with SHA256 checksums for integrity validation.
+    """
+
+    @staticmethod
+    def compute_checksum(payload: dict[str, Any]) -> str:
+        """
+        Compute SHA256 checksum for payload.
+
+        Args:
+            payload: State payload
+
+        Returns:
+            Hex-encoded SHA256 checksum
+        """
+        payload_json = json.dumps(payload, sort_keys=True)
+        return hashlib.sha256(payload_json.encode()).hexdigest()
+
+    def serialize(
+        self,
+        state: dict[str, Any],
+        processor_id: str,
+        pod_name: str,
+        restart_count: int = 0,
+    ) -> str:
+        """
+        Serialize state to JSON string with envelope.
+
+        Args:
+            state: State dictionary to serialize
+            processor_id: Processor ID for metadata
+            pod_name: Pod name for metadata
+            restart_count: Current restart count
+
+        Returns:
+            JSON string with state envelope
+        """
+        envelope = StateEnvelope(
+            payload=state,
+            metadata={
+                "timestamp": time.time(),
+                "timestamp_iso": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                "processor_id": processor_id,
+                "pod_name": pod_name,
+                "restart_count": restart_count,
+            },
+            checksum=self.compute_checksum(state),
+        )
+
+        return json.dumps(envelope.to_dict(), indent=2)
+
+    def deserialize(self, data: str) -> dict[str, Any]:
+        """
+        Deserialize state from JSON string.
+
+        Args:
+            data: JSON string with state envelope
+
+        Returns:
+            State payload dictionary
+
+        Raises:
+            DoryStateError: If deserialization or validation fails
+        """
+        try:
+            envelope_dict = json.loads(data)
+        except json.JSONDecodeError as e:
+            raise DoryStateError(f"Invalid JSON in state data: {e}", cause=e)
+
+        try:
+            envelope = StateEnvelope.from_dict(envelope_dict)
+        except KeyError as e:
+            raise DoryStateError(f"Missing field in state envelope: {e}", cause=e)
+
+        # Validate checksum
+        expected_checksum = self.compute_checksum(envelope.payload)
+        if envelope.checksum != expected_checksum:
+            raise DoryStateError(
+                f"State checksum mismatch: expected {expected_checksum}, "
+                f"got {envelope.checksum}"
+            )
+
+        return envelope.payload
+

dory/migration/state_manager.py

@@ -0,0 +1,286 @@
+"""
+StateManager - High-level state management for migrations.
+
+Provides unified interface for state operations across different
+storage backends (ConfigMap, S3, local file).
+"""
+
+import logging
+import os
+from pathlib import Path
+from typing import Any, TYPE_CHECKING
+
+from dory.types import StateBackend
+from dory.migration.serialization import StateSerializer
+from dory.migration.configmap import ConfigMapStore
+from dory.migration.s3_store import S3Store, S3Config
+from dory.utils.errors import DoryStateError
+
+if TYPE_CHECKING:
+    from dory.config.schema import DoryConfig
+
+logger = logging.getLogger(__name__)
+
+
+class StateManager:
+    """
+    High-level state management for processor migrations.
+
+    Supports multiple backends:
+    - ConfigMap: Kubernetes ConfigMap (default, <1MB)
+    - S3: AWS S3 (for multi-cluster)
+    - Local/PVC: Local file (tries /data first, falls back to /tmp)
+    """
+
+    LOCAL_STATE_PATH = "/data/dory-state.json"
+    LOCAL_STATE_PATH_FALLBACK = "/tmp/dory-state.json"
+
+    def __init__(
+        self,
+        backend: str | StateBackend = StateBackend.CONFIGMAP,
+        config: "DoryConfig | None" = None,
+    ):
+        """
+        Initialize state manager.
+
+        Args:
+            backend: Storage backend to use
+            config: SDK configuration
+        """
+        if isinstance(backend, str):
+            backend = StateBackend(backend)
+
+        self._backend = backend
+        self._config = config
+        self._serializer = StateSerializer()
+        self._configmap_store: ConfigMapStore | None = None
+        self._s3_store: S3Store | None = None
+
+        # Get namespace from environment
+        self._namespace = os.environ.get("POD_NAMESPACE", "default")
+        self._pod_name = os.environ.get("POD_NAME", "unknown")
+
+    async def save_state(
+        self,
+        processor_id: str,
+        state: dict[str, Any],
+        restart_count: int = 0,
+    ) -> None:
+        """
+        Save processor state.
+
+        Args:
+            processor_id: Processor ID
+            state: State dictionary to save
+            restart_count: Current restart count for metadata
+
+        Raises:
+            DoryStateError: If save fails
+        """
+        logger.debug(f"Saving state for processor {processor_id}")
+
+        # Serialize state
+        state_json = self._serializer.serialize(
+            state=state,
+            processor_id=processor_id,
+            pod_name=self._pod_name,
+            restart_count=restart_count,
+        )
+
+        # Save to backend (PVC uses the same local file path as Local)
+        if self._backend == StateBackend.CONFIGMAP:
+            await self._save_to_configmap(processor_id, state_json)
+        elif self._backend in (StateBackend.LOCAL, StateBackend.PVC):
+            await self._save_to_local(processor_id, state_json)
+        elif self._backend == StateBackend.S3:
+            await self._save_to_s3(processor_id, state_json)
+        else:
+            raise DoryStateError(f"Unsupported state backend: {self._backend}")
+
+        logger.info(f"State saved for processor {processor_id}")
+
+    async def load_state(self, processor_id: str) -> dict[str, Any] | None:
+        """
+        Load processor state.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            State dictionary, or None if no state found
+
+        Raises:
+            DoryStateError: If load fails
+        """
+        logger.debug(f"Loading state for processor {processor_id}")
+
+        # Load from backend
+        state_json: str | None = None
+
+        if self._backend == StateBackend.CONFIGMAP:
+            state_json = await self._load_from_configmap(processor_id)
+        elif self._backend in (StateBackend.LOCAL, StateBackend.PVC):
+            state_json = await self._load_from_local(processor_id)
+        elif self._backend == StateBackend.S3:
+            state_json = await self._load_from_s3(processor_id)
+        else:
+            raise DoryStateError(f"Unsupported state backend: {self._backend}")
+
+        if state_json is None:
+            logger.debug(f"No state found for processor {processor_id}")
+            return None
+
+        # Deserialize state
+        state = self._serializer.deserialize(state_json)
+        logger.info(f"State loaded for processor {processor_id}")
+        return state
+
+    async def delete_state(self, processor_id: str) -> bool:
+        """
+        Delete processor state (golden image reset).
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            True if state was deleted
+
+        Raises:
+            DoryStateError: If delete fails
+        """
+        logger.debug(f"Deleting state for processor {processor_id}")
+
+        if self._backend == StateBackend.CONFIGMAP:
+            return await self._delete_from_configmap(processor_id)
+        elif self._backend in (StateBackend.LOCAL, StateBackend.PVC):
+            return await self._delete_from_local(processor_id)
+        elif self._backend == StateBackend.S3:
+            return await self._delete_from_s3(processor_id)
+        else:
+            raise DoryStateError(f"Unsupported state backend: {self._backend}")
+
+    # =========================================================================
+    # ConfigMap Backend
+    # =========================================================================
+
+    def _get_configmap_store(self) -> ConfigMapStore:
+        """Get or create ConfigMap store instance."""
+        if self._configmap_store is None:
+            self._configmap_store = ConfigMapStore(namespace=self._namespace)
+        return self._configmap_store
+
+    async def _save_to_configmap(self, processor_id: str, state_json: str) -> None:
+        """Save state to Kubernetes ConfigMap."""
+        await self._get_configmap_store().save(processor_id, state_json)
+
+    async def _load_from_configmap(self, processor_id: str) -> str | None:
+        """Load state from Kubernetes ConfigMap."""
+        return await self._get_configmap_store().load(processor_id)
+
+    async def _delete_from_configmap(self, processor_id: str) -> bool:
+        """Delete state ConfigMap."""
+        return await self._get_configmap_store().delete(processor_id)
+
+    # =========================================================================
+    # Local File Backend (also used by PVC — tries /data first, then /tmp)
+    # =========================================================================
+
+    def _get_local_path(self, processor_id: str) -> Path:
+        """Get local file path for state."""
+        # Try /data first (works when PVC is mounted), fall back to /tmp
+        base_path = Path(self.LOCAL_STATE_PATH).parent
+        if not base_path.exists():
+            base_path = Path(self.LOCAL_STATE_PATH_FALLBACK).parent
+
+        return base_path / f"dory-state-{processor_id}.json"
+
+    async def _save_to_local(self, processor_id: str, state_json: str) -> None:
+        """Save state to local file."""
+        path = self._get_local_path(processor_id)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            path.write_text(state_json)
+            logger.debug(f"State saved to local file: {path}")
+        except Exception as e:
+            raise DoryStateError(f"Failed to save state to {path}: {e}", cause=e)
+
+    async def _load_from_local(self, processor_id: str) -> str | None:
+        """Load state from local file."""
+        path = self._get_local_path(processor_id)
+
+        if not path.exists():
+            return None
+
+        try:
+            return path.read_text()
+        except Exception as e:
+            raise DoryStateError(f"Failed to load state from {path}: {e}", cause=e)
+
+    async def _delete_from_local(self, processor_id: str) -> bool:
+        """Delete local state file."""
+        path = self._get_local_path(processor_id)
+
+        if not path.exists():
+            return False
+
+        try:
+            path.unlink()
+            logger.debug(f"State file deleted: {path}")
+            return True
+        except Exception as e:
+            raise DoryStateError(f"Failed to delete state file {path}: {e}", cause=e)
+
+    # =========================================================================
+    # S3 Backend
+    # =========================================================================
+
+    def _get_s3_store(self) -> S3Store:
+        """Get or create S3 store instance."""
+        if self._s3_store is None:
+            # Try to get S3 config from DoryConfig if available
+            s3_config = None
+            if self._config and hasattr(self._config, "s3_config"):
+                s3_config = self._config.s3_config
+
+            self._s3_store = S3Store(config=s3_config)
+
+        return self._s3_store
+
+    async def _save_to_s3(self, processor_id: str, state_json: str) -> None:
+        """Save state to S3 with offline buffering support."""
+        store = self._get_s3_store()
+        await store.save(
+            processor_id,
+            state_json,
+            metadata={
+                "pod-name": self._pod_name,
+                "namespace": self._namespace,
+            },
+        )
+
+    async def _load_from_s3(self, processor_id: str) -> str | None:
+        """Load state from S3 (falls back to local buffer if unavailable)."""
+        store = self._get_s3_store()
+        return await store.load(processor_id)
+
+    async def _delete_from_s3(self, processor_id: str) -> bool:
+        """Delete state from S3."""
+        store = self._get_s3_store()
+        return await store.delete(processor_id)
+
+    async def sync_s3_buffer(self) -> int:
+        """
+        Sync locally buffered states to S3.
+
+        Call this periodically on edge nodes to upload states
+        that were buffered during connectivity issues.
+
+        Returns:
+            Number of states synced
+        """
+        if self._backend != StateBackend.S3:
+            return 0
+
+        store = self._get_s3_store()
+        return await store.sync_buffer()