dory-sdk 2.1.0__py3-none-any.whl

dory/migration/configmap.py

@@ -0,0 +1,260 @@
+"""
+ConfigMap storage backend for state persistence.
+
+Uses Kubernetes ConfigMaps to store processor state during migrations.
+"""
+
+import logging
+import time
+from typing import Any
+
+from dory.utils.errors import DoryK8sError, DoryStateError
+
+logger = logging.getLogger(__name__)
+
+# Optional kubernetes import - gracefully handle if not available
+try:
+    from kubernetes import client, config
+    from kubernetes.client.rest import ApiException
+    K8S_AVAILABLE = True
+except ImportError:
+    K8S_AVAILABLE = False
+    client = None
+    config = None
+    ApiException = Exception
+
+
+class ConfigMapStore:
+    """
+    Store and retrieve state from Kubernetes ConfigMaps.
+
+    ConfigMap naming convention: dory-state-{processor_id}
+    TTL: Auto-cleanup after 1 hour if not claimed.
+    """
+
+    STATE_CONFIGMAP_PREFIX = "dory-state-"
+    STATE_KEY = "state"
+    TTL_ANNOTATION = "dory.io/state-ttl"
+    CREATED_ANNOTATION = "dory.io/created-timestamp"
+    OWNER_LABEL = "dory.io/state-owner"
+    DEFAULT_TTL_SECONDS = 3600  # 1 hour
+
+    def __init__(self, namespace: str | None = None):
+        """
+        Initialize ConfigMap store.
+
+        Args:
+            namespace: Kubernetes namespace (defaults to current pod's namespace)
+        """
+        self._namespace = namespace
+        self._api: Any = None
+        self._initialized = False
+
+    def _ensure_initialized(self) -> None:
+        """Initialize Kubernetes client if not already done."""
+        if self._initialized:
+            return
+
+        if not K8S_AVAILABLE:
+            raise DoryK8sError(
+                "Kubernetes client not available. "
+                "Install with: pip install kubernetes"
+            )
+
+        try:
+            # Try in-cluster config first
+            config.load_incluster_config()
+            logger.debug("Using in-cluster Kubernetes config")
+        except config.ConfigException:
+            try:
+                # Fall back to kubeconfig
+                config.load_kube_config()
+                logger.debug("Using kubeconfig")
+            except config.ConfigException as e:
+                raise DoryK8sError(f"Failed to load Kubernetes config: {e}", cause=e)
+
+        self._api = client.CoreV1Api()
+
+        # Get namespace from pod environment if not specified
+        if not self._namespace:
+            import os
+            self._namespace = os.environ.get("POD_NAMESPACE", "default")
+
+        self._initialized = True
+
+    def _configmap_name(self, processor_id: str) -> str:
+        """Generate ConfigMap name for processor."""
+        return f"{self.STATE_CONFIGMAP_PREFIX}{processor_id}"
+
+    async def save(
+        self,
+        processor_id: str,
+        state_json: str,
+        ttl_seconds: int | None = None,
+    ) -> None:
+        """
+        Save state to ConfigMap.
+
+        Args:
+            processor_id: Processor ID
+            state_json: JSON-serialized state
+            ttl_seconds: TTL for auto-cleanup (default 1 hour)
+
+        Raises:
+            DoryK8sError: If ConfigMap operation fails
+        """
+        self._ensure_initialized()
+
+        cm_name = self._configmap_name(processor_id)
+        ttl = ttl_seconds or self.DEFAULT_TTL_SECONDS
+
+        configmap = client.V1ConfigMap(
+            metadata=client.V1ObjectMeta(
+                name=cm_name,
+                namespace=self._namespace,
+                labels={
+                    self.OWNER_LABEL: "true",
+                },
+                annotations={
+                    self.TTL_ANNOTATION: str(ttl),
+                    self.CREATED_ANNOTATION: time.strftime(
+                        "%Y-%m-%dT%H:%M:%SZ", time.gmtime()
+                    ),
+                },
+            ),
+            data={
+                self.STATE_KEY: state_json,
+            },
+        )
+
+        try:
+            # Try to create first
+            self._api.create_namespaced_config_map(
+                namespace=self._namespace,
+                body=configmap,
+            )
+            logger.debug(f"Created state ConfigMap: {cm_name}")
+
+        except ApiException as e:
+            if e.status == 409:
+                # Already exists, update it
+                try:
+                    self._api.replace_namespaced_config_map(
+                        name=cm_name,
+                        namespace=self._namespace,
+                        body=configmap,
+                    )
+                    logger.debug(f"Updated state ConfigMap: {cm_name}")
+                except ApiException as e2:
+                    raise DoryK8sError(
+                        f"Failed to update ConfigMap {cm_name}: {e2}",
+                        cause=e2,
+                    )
+            else:
+                raise DoryK8sError(
+                    f"Failed to create ConfigMap {cm_name}: {e}",
+                    cause=e,
+                )
+
+    async def load(self, processor_id: str) -> str | None:
+        """
+        Load state from ConfigMap.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            JSON-serialized state, or None if not found
+
+        Raises:
+            DoryK8sError: If ConfigMap operation fails
+        """
+        self._ensure_initialized()
+
+        cm_name = self._configmap_name(processor_id)
+
+        try:
+            configmap = self._api.read_namespaced_config_map(
+                name=cm_name,
+                namespace=self._namespace,
+            )
+
+            state_json = configmap.data.get(self.STATE_KEY)
+            if state_json:
+                logger.debug(f"Loaded state from ConfigMap: {cm_name}")
+                return state_json
+            else:
+                logger.warning(f"ConfigMap {cm_name} exists but has no state data")
+                return None
+
+        except ApiException as e:
+            if e.status == 404:
+                logger.debug(f"State ConfigMap not found: {cm_name}")
+                return None
+            raise DoryK8sError(
+                f"Failed to read ConfigMap {cm_name}: {e}",
+                cause=e,
+            )
+
+    async def delete(self, processor_id: str) -> bool:
+        """
+        Delete state ConfigMap.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            True if deleted, False if not found
+
+        Raises:
+            DoryK8sError: If ConfigMap operation fails
+        """
+        self._ensure_initialized()
+
+        cm_name = self._configmap_name(processor_id)
+
+        try:
+            self._api.delete_namespaced_config_map(
+                name=cm_name,
+                namespace=self._namespace,
+            )
+            logger.debug(f"Deleted state ConfigMap: {cm_name}")
+            return True
+
+        except ApiException as e:
+            if e.status == 404:
+                logger.debug(f"State ConfigMap not found for deletion: {cm_name}")
+                return False
+            raise DoryK8sError(
+                f"Failed to delete ConfigMap {cm_name}: {e}",
+                cause=e,
+            )
+
+    async def exists(self, processor_id: str) -> bool:
+        """
+        Check if state ConfigMap exists.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            True if ConfigMap exists
+        """
+        self._ensure_initialized()
+
+        cm_name = self._configmap_name(processor_id)
+
+        try:
+            self._api.read_namespaced_config_map(
+                name=cm_name,
+                namespace=self._namespace,
+            )
+            return True
+
+        except ApiException as e:
+            if e.status == 404:
+                return False
+            raise DoryK8sError(
+                f"Failed to check ConfigMap {cm_name}: {e}",
+                cause=e,
+            )

dory/migration/serialization.py

@@ -0,0 +1,167 @@
+"""
+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
+
+    def deserialize_with_metadata(self, data: str) -> StateEnvelope:
+        """
+        Deserialize state with full envelope including metadata.
+
+        Args:
+            data: JSON string with state envelope
+
+        Returns:
+            StateEnvelope with payload and metadata
+
+        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