dory-sdk 2.1.0__py3-none-any.whl → 2.1.4__py3-none-any.whl

dory/migration/s3_store.py

@@ -0,0 +1,656 @@
+"""
+S3 storage backend for state persistence.
+
+Provides S3-based state storage with local buffering for edge nodes
+that may have intermittent connectivity.
+
+Features:
+- S3 upload/download with retry logic
+- Local SQLite buffer for offline scenarios
+- Multiple credential options (IAM role, env vars, STS)
+- Automatic sync when connectivity is restored
+"""
+
+import asyncio
+import json
+import logging
+import os
+import sqlite3
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from dory.utils.errors import DoryStateError
+
+logger = logging.getLogger(__name__)
+
+# Optional boto3 import - gracefully handle if not available
+try:
+    import boto3
+    from botocore.exceptions import ClientError, NoCredentialsError, BotoCoreError
+    BOTO3_AVAILABLE = True
+except ImportError:
+    BOTO3_AVAILABLE = False
+    boto3 = None
+    ClientError = Exception
+    NoCredentialsError = Exception
+    BotoCoreError = Exception
+
+
+@dataclass
+class S3Config:
+    """Configuration for S3 state backend."""
+
+    bucket: str
+    prefix: str = "dory-state"
+    region: str | None = None
+    endpoint_url: str | None = None  # For S3-compatible services (MinIO, LocalStack)
+
+    # Credential options
+    access_key_id: str | None = None
+    secret_access_key: str | None = None
+    session_token: str | None = None  # For STS temporary credentials
+    role_arn: str | None = None  # For assuming a role
+    credential_broker_url: str | None = None  # For edge credential broker
+
+    # Offline buffering
+    enable_offline_buffer: bool = True
+    buffer_path: str = "/data/dory-state-buffer.db"
+    buffer_path_fallback: str = "/tmp/dory-state-buffer.db"
+    max_buffer_age_seconds: int = 86400  # 24 hours
+
+    # Retry settings
+    max_retries: int = 3
+    retry_delay_seconds: float = 1.0
+    retry_backoff_multiplier: float = 2.0
+
+    @classmethod
+    def from_env(cls) -> "S3Config":
+        """Create config from environment variables."""
+        bucket = os.environ.get("DORY_S3_BUCKET")
+        if not bucket:
+            raise DoryStateError(
+                "DORY_S3_BUCKET environment variable is required for S3 backend"
+            )
+
+        return cls(
+            bucket=bucket,
+            prefix=os.environ.get("DORY_S3_PREFIX", "dory-state"),
+            region=os.environ.get("DORY_S3_REGION", os.environ.get("AWS_REGION")),
+            endpoint_url=os.environ.get("DORY_S3_ENDPOINT_URL"),
+            access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
+            secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
+            session_token=os.environ.get("AWS_SESSION_TOKEN"),
+            role_arn=os.environ.get("DORY_S3_ROLE_ARN"),
+            credential_broker_url=os.environ.get("DORY_CREDENTIAL_BROKER_URL"),
+            enable_offline_buffer=os.environ.get(
+                "DORY_S3_ENABLE_OFFLINE_BUFFER", "true"
+            ).lower() == "true",
+            buffer_path=os.environ.get(
+                "DORY_S3_BUFFER_PATH", "/data/dory-state-buffer.db"
+            ),
+        )
+
+
+class OfflineBuffer:
+    """
+    SQLite-based local buffer for offline state persistence.
+
+    Stores state locally when S3 is unreachable and syncs when
+    connectivity is restored.
+    """
+
+    def __init__(self, db_path: str, fallback_path: str):
+        """
+        Initialize offline buffer.
+
+        Args:
+            db_path: Primary path for SQLite database
+            fallback_path: Fallback path if primary is not writable
+        """
+        self._db_path = self._resolve_path(db_path, fallback_path)
+        self._conn: sqlite3.Connection | None = None
+        self._initialized = False
+
+    def _resolve_path(self, primary: str, fallback: str) -> str:
+        """Resolve which path to use for the database."""
+        primary_dir = Path(primary).parent
+        if primary_dir.exists() and os.access(primary_dir, os.W_OK):
+            return primary
+
+        fallback_dir = Path(fallback).parent
+        fallback_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Using fallback buffer path: {fallback}")
+        return fallback
+
+    def _ensure_initialized(self) -> None:
+        """Initialize database if not already done."""
+        if self._initialized:
+            return
+
+        self._conn = sqlite3.connect(self._db_path)
+        self._conn.execute("""
+            CREATE TABLE IF NOT EXISTS state_buffer (
+                processor_id TEXT PRIMARY KEY,
+                state_json TEXT NOT NULL,
+                created_at REAL NOT NULL,
+                synced_at REAL,
+                sync_attempts INTEGER DEFAULT 0
+            )
+        """)
+        self._conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_synced
+            ON state_buffer(synced_at)
+        """)
+        self._conn.commit()
+        self._initialized = True
+        logger.debug(f"Offline buffer initialized at {self._db_path}")
+
+    def save(self, processor_id: str, state_json: str) -> None:
+        """Save state to local buffer."""
+        self._ensure_initialized()
+
+        self._conn.execute("""
+            INSERT OR REPLACE INTO state_buffer
+            (processor_id, state_json, created_at, synced_at, sync_attempts)
+            VALUES (?, ?, ?, NULL, 0)
+        """, (processor_id, state_json, time.time()))
+        self._conn.commit()
+        logger.debug(f"State buffered locally for {processor_id}")
+
+    def load(self, processor_id: str) -> str | None:
+        """Load state from local buffer."""
+        self._ensure_initialized()
+
+        cursor = self._conn.execute("""
+            SELECT state_json FROM state_buffer
+            WHERE processor_id = ?
+        """, (processor_id,))
+        row = cursor.fetchone()
+        return row[0] if row else None
+
+    def mark_synced(self, processor_id: str) -> None:
+        """Mark state as synced to S3."""
+        self._ensure_initialized()
+
+        self._conn.execute("""
+            UPDATE state_buffer
+            SET synced_at = ?
+            WHERE processor_id = ?
+        """, (time.time(), processor_id))
+        self._conn.commit()
+
+    def get_unsynced(self) -> list[tuple[str, str]]:
+        """Get all unsynced states."""
+        self._ensure_initialized()
+
+        cursor = self._conn.execute("""
+            SELECT processor_id, state_json
+            FROM state_buffer
+            WHERE synced_at IS NULL
+            ORDER BY created_at ASC
+        """)
+        return cursor.fetchall()
+
+    def increment_sync_attempts(self, processor_id: str) -> None:
+        """Increment sync attempt counter."""
+        self._ensure_initialized()
+
+        self._conn.execute("""
+            UPDATE state_buffer
+            SET sync_attempts = sync_attempts + 1
+            WHERE processor_id = ?
+        """, (processor_id,))
+        self._conn.commit()
+
+    def delete(self, processor_id: str) -> bool:
+        """Delete state from buffer."""
+        self._ensure_initialized()
+
+        cursor = self._conn.execute("""
+            DELETE FROM state_buffer WHERE processor_id = ?
+        """, (processor_id,))
+        self._conn.commit()
+        return cursor.rowcount > 0
+
+    def cleanup_old(self, max_age_seconds: int) -> int:
+        """Remove entries older than max_age_seconds."""
+        self._ensure_initialized()
+
+        cutoff = time.time() - max_age_seconds
+        cursor = self._conn.execute("""
+            DELETE FROM state_buffer
+            WHERE created_at < ? AND synced_at IS NOT NULL
+        """, (cutoff,))
+        self._conn.commit()
+        return cursor.rowcount
+
+    def close(self) -> None:
+        """Close database connection."""
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+            self._initialized = False
+
+
+class S3Store:
+    """
+    Store and retrieve state from AWS S3.
+
+    Supports offline buffering for edge nodes with intermittent connectivity.
+
+    S3 key format: {prefix}/{processor_id}/state.json
+    """
+
+    def __init__(self, config: S3Config | None = None):
+        """
+        Initialize S3 store.
+
+        Args:
+            config: S3 configuration (defaults to from_env())
+        """
+        if not BOTO3_AVAILABLE:
+            raise DoryStateError(
+                "boto3 is required for S3 backend. "
+                "Install with: pip install boto3"
+            )
+
+        self._config = config or S3Config.from_env()
+        self._client: Any = None
+        self._initialized = False
+
+        # Offline buffer
+        self._buffer: OfflineBuffer | None = None
+        if self._config.enable_offline_buffer:
+            self._buffer = OfflineBuffer(
+                self._config.buffer_path,
+                self._config.buffer_path_fallback,
+            )
+
+        # Background sync task
+        self._sync_task: asyncio.Task | None = None
+
+    def _ensure_initialized(self) -> None:
+        """Initialize S3 client if not already done."""
+        if self._initialized:
+            return
+
+        try:
+            session_kwargs = {}
+            client_kwargs = {}
+
+            # Region
+            if self._config.region:
+                session_kwargs["region_name"] = self._config.region
+
+            # Explicit credentials
+            if self._config.access_key_id and self._config.secret_access_key:
+                session_kwargs["aws_access_key_id"] = self._config.access_key_id
+                session_kwargs["aws_secret_access_key"] = self._config.secret_access_key
+                if self._config.session_token:
+                    session_kwargs["aws_session_token"] = self._config.session_token
+
+            # Custom endpoint (MinIO, LocalStack)
+            if self._config.endpoint_url:
+                client_kwargs["endpoint_url"] = self._config.endpoint_url
+
+            # Create session and client
+            session = boto3.Session(**session_kwargs)
+
+            # Assume role if specified
+            if self._config.role_arn:
+                sts = session.client("sts")
+                assumed = sts.assume_role(
+                    RoleArn=self._config.role_arn,
+                    RoleSessionName="dory-state-manager",
+                )
+                credentials = assumed["Credentials"]
+                session = boto3.Session(
+                    aws_access_key_id=credentials["AccessKeyId"],
+                    aws_secret_access_key=credentials["SecretAccessKey"],
+                    aws_session_token=credentials["SessionToken"],
+                    region_name=self._config.region,
+                )
+
+            self._client = session.client("s3", **client_kwargs)
+            self._initialized = True
+            logger.debug(f"S3 client initialized for bucket {self._config.bucket}")
+
+        except NoCredentialsError as e:
+            raise DoryStateError(
+                "No AWS credentials found. Configure via environment variables, "
+                "IAM role, or credential broker.",
+                cause=e,
+            )
+        except Exception as e:
+            raise DoryStateError(f"Failed to initialize S3 client: {e}", cause=e)
+
+    def _s3_key(self, processor_id: str) -> str:
+        """Generate S3 key for processor state."""
+        return f"{self._config.prefix}/{processor_id}/state.json"
+
+    async def _retry_operation(
+        self,
+        operation: str,
+        func: Any,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """Execute operation with retry logic."""
+        last_error = None
+        delay = self._config.retry_delay_seconds
+
+        for attempt in range(1, self._config.max_retries + 1):
+            try:
+                # Run synchronous boto3 call in executor
+                loop = asyncio.get_event_loop()
+                result = await loop.run_in_executor(None, lambda: func(*args, **kwargs))
+                return result
+
+            except (ClientError, BotoCoreError) as e:
+                last_error = e
+                if attempt < self._config.max_retries:
+                    logger.warning(
+                        f"S3 {operation} attempt {attempt} failed: {e}. "
+                        f"Retrying in {delay:.1f}s..."
+                    )
+                    await asyncio.sleep(delay)
+                    delay *= self._config.retry_backoff_multiplier
+
+        raise DoryStateError(
+            f"S3 {operation} failed after {self._config.max_retries} attempts: {last_error}",
+            cause=last_error,
+        )
+
+    async def save(
+        self,
+        processor_id: str,
+        state_json: str,
+        metadata: dict[str, str] | None = None,
+    ) -> None:
+        """
+        Save state to S3.
+
+        If S3 is unreachable and offline buffering is enabled,
+        state is saved to local buffer for later sync.
+
+        Args:
+            processor_id: Processor ID
+            state_json: JSON-serialized state
+            metadata: Optional metadata to store with object
+
+        Raises:
+            DoryStateError: If save fails and no buffer available
+        """
+        self._ensure_initialized()
+
+        s3_key = self._s3_key(processor_id)
+        s3_metadata = metadata or {}
+        s3_metadata["processor-id"] = processor_id
+        s3_metadata["saved-at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+
+        try:
+            await self._retry_operation(
+                "put_object",
+                self._client.put_object,
+                Bucket=self._config.bucket,
+                Key=s3_key,
+                Body=state_json.encode("utf-8"),
+                ContentType="application/json",
+                Metadata=s3_metadata,
+            )
+            logger.debug(f"State saved to S3: s3://{self._config.bucket}/{s3_key}")
+
+            # If buffer exists, mark as synced
+            if self._buffer:
+                self._buffer.save(processor_id, state_json)
+                self._buffer.mark_synced(processor_id)
+
+        except DoryStateError:
+            # S3 failed - try to buffer locally
+            if self._buffer:
+                self._buffer.save(processor_id, state_json)
+                logger.warning(
+                    f"S3 unavailable, state buffered locally for {processor_id}"
+                )
+            else:
+                raise
+
+    async def load(self, processor_id: str) -> str | None:
+        """
+        Load state from S3.
+
+        Falls back to local buffer if S3 is unreachable.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            JSON-serialized state, or None if not found
+
+        Raises:
+            DoryStateError: If load fails
+        """
+        self._ensure_initialized()
+
+        s3_key = self._s3_key(processor_id)
+
+        try:
+            response = await self._retry_operation(
+                "get_object",
+                self._client.get_object,
+                Bucket=self._config.bucket,
+                Key=s3_key,
+            )
+
+            # Read body in executor
+            loop = asyncio.get_event_loop()
+            body = await loop.run_in_executor(
+                None,
+                lambda: response["Body"].read().decode("utf-8"),
+            )
+
+            logger.debug(f"State loaded from S3: s3://{self._config.bucket}/{s3_key}")
+            return body
+
+        except DoryStateError as e:
+            # Check if it's a 404
+            if "NoSuchKey" in str(e) or "404" in str(e):
+                logger.debug(f"State not found in S3: {s3_key}")
+
+                # Try local buffer as fallback
+                if self._buffer:
+                    buffered = self._buffer.load(processor_id)
+                    if buffered:
+                        logger.info(
+                            f"Using buffered state for {processor_id} (not yet synced to S3)"
+                        )
+                        return buffered
+
+                return None
+
+            # S3 error - try local buffer
+            if self._buffer:
+                buffered = self._buffer.load(processor_id)
+                if buffered:
+                    logger.warning(
+                        f"S3 unavailable, using buffered state for {processor_id}"
+                    )
+                    return buffered
+
+            raise
+
+        except ClientError as e:
+            if e.response.get("Error", {}).get("Code") == "NoSuchKey":
+                logger.debug(f"State not found in S3: {s3_key}")
+                return None
+            raise DoryStateError(f"Failed to load state from S3: {e}", cause=e)
+
+    async def delete(self, processor_id: str) -> bool:
+        """
+        Delete state from S3.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            True if deleted, False if not found
+
+        Raises:
+            DoryStateError: If delete fails
+        """
+        self._ensure_initialized()
+
+        s3_key = self._s3_key(processor_id)
+
+        try:
+            # Check if exists first
+            try:
+                await self._retry_operation(
+                    "head_object",
+                    self._client.head_object,
+                    Bucket=self._config.bucket,
+                    Key=s3_key,
+                )
+            except DoryStateError:
+                # Not found
+                if self._buffer:
+                    self._buffer.delete(processor_id)
+                return False
+
+            # Delete from S3
+            await self._retry_operation(
+                "delete_object",
+                self._client.delete_object,
+                Bucket=self._config.bucket,
+                Key=s3_key,
+            )
+            logger.debug(f"State deleted from S3: s3://{self._config.bucket}/{s3_key}")
+
+            # Delete from buffer too
+            if self._buffer:
+                self._buffer.delete(processor_id)
+
+            return True
+
+        except ClientError as e:
+            if e.response.get("Error", {}).get("Code") == "NoSuchKey":
+                return False
+            raise DoryStateError(f"Failed to delete state from S3: {e}", cause=e)
+
+    async def exists(self, processor_id: str) -> bool:
+        """
+        Check if state exists in S3.
+
+        Args:
+            processor_id: Processor ID
+
+        Returns:
+            True if state exists
+        """
+        self._ensure_initialized()
+
+        s3_key = self._s3_key(processor_id)
+
+        try:
+            await self._retry_operation(
+                "head_object",
+                self._client.head_object,
+                Bucket=self._config.bucket,
+                Key=s3_key,
+            )
+            return True
+        except (DoryStateError, ClientError):
+            return False
+
+    async def sync_buffer(self) -> int:
+        """
+        Sync buffered states to S3.
+
+        Call periodically to upload states that were buffered
+        during connectivity issues.
+
+        Returns:
+            Number of states synced
+        """
+        if not self._buffer:
+            return 0
+
+        unsynced = self._buffer.get_unsynced()
+        synced_count = 0
+
+        for processor_id, state_json in unsynced:
+            try:
+                s3_key = self._s3_key(processor_id)
+
+                await self._retry_operation(
+                    "put_object",
+                    self._client.put_object,
+                    Bucket=self._config.bucket,
+                    Key=s3_key,
+                    Body=state_json.encode("utf-8"),
+                    ContentType="application/json",
+                    Metadata={
+                        "processor-id": processor_id,
+                        "synced-from-buffer": "true",
+                        "synced-at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                    },
+                )
+
+                self._buffer.mark_synced(processor_id)
+                synced_count += 1
+                logger.info(f"Synced buffered state for {processor_id} to S3")
+
+            except DoryStateError as e:
+                self._buffer.increment_sync_attempts(processor_id)
+                logger.warning(f"Failed to sync buffered state for {processor_id}: {e}")
+
+        # Cleanup old synced entries
+        if synced_count > 0:
+            cleaned = self._buffer.cleanup_old(self._config.max_buffer_age_seconds)
+            if cleaned > 0:
+                logger.debug(f"Cleaned up {cleaned} old buffer entries")
+
+        return synced_count
+
+    async def start_background_sync(self, interval_seconds: float = 60.0) -> None:
+        """
+        Start background sync task.
+
+        Args:
+            interval_seconds: Interval between sync attempts
+        """
+        if self._sync_task is not None:
+            return
+
+        async def sync_loop():
+            while True:
+                try:
+                    await asyncio.sleep(interval_seconds)
+                    synced = await self.sync_buffer()
+                    if synced > 0:
+                        logger.debug(f"Background sync: {synced} states synced")
+                except asyncio.CancelledError:
+                    break
+                except Exception as e:
+                    logger.error(f"Background sync error: {e}")
+
+        self._sync_task = asyncio.create_task(sync_loop())
+        logger.info(f"Started background S3 sync (interval: {interval_seconds}s)")
+
+    async def stop_background_sync(self) -> None:
+        """Stop background sync task."""
+        if self._sync_task:
+            self._sync_task.cancel()
+            try:
+                await self._sync_task
+            except asyncio.CancelledError:
+                pass
+            self._sync_task = None
+            logger.info("Stopped background S3 sync")
+
+    def close(self) -> None:
+        """Close resources."""
+        if self._buffer:
+            self._buffer.close()