kailash 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

kailash/resources/health.py

@@ -0,0 +1,319 @@
+"""
+Health Check System - Monitor and maintain resource health.
+
+This module provides health checking functionality for resources in the registry,
+enabling automatic recovery and circuit breaker patterns.
+"""
+
+import asyncio
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, Optional, Protocol, Union
+
+
+class HealthCheckProtocol(Protocol):
+    """Protocol for health check callables."""
+
+    async def __call__(self, resource: Any) -> Union[bool, "HealthStatus"]:
+        """
+        Check if a resource is healthy.
+
+        Args:
+            resource: The resource to check
+
+        Returns:
+            bool or HealthStatus indicating health
+        """
+        ...
+
+
+# Type alias for health checks
+HealthCheck = HealthCheckProtocol
+
+
+class HealthState(Enum):
+    """Health states for resources."""
+
+    HEALTHY = "healthy"
+    DEGRADED = "degraded"
+    UNHEALTHY = "unhealthy"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class HealthStatus:
+    """
+    Detailed health status for a resource.
+
+    This provides more information than a simple boolean,
+    allowing for degraded states and detailed diagnostics.
+    """
+
+    state: HealthState
+    message: Optional[str] = None
+    details: Optional[Dict[str, Any]] = None
+    timestamp: datetime = None
+
+    def __post_init__(self):
+        if self.timestamp is None:
+            self.timestamp = datetime.now()
+
+    @property
+    def is_healthy(self) -> bool:
+        """Check if status indicates health."""
+        return self.state in (HealthState.HEALTHY, HealthState.DEGRADED)
+
+    @classmethod
+    def healthy(cls, message: str = "Resource is healthy") -> "HealthStatus":
+        """Create a healthy status."""
+        return cls(HealthState.HEALTHY, message)
+
+    @classmethod
+    def unhealthy(cls, message: str, details: Dict[str, Any] = None) -> "HealthStatus":
+        """Create an unhealthy status."""
+        return cls(HealthState.UNHEALTHY, message, details)
+
+    @classmethod
+    def degraded(cls, message: str, details: Dict[str, Any] = None) -> "HealthStatus":
+        """Create a degraded status."""
+        return cls(HealthState.DEGRADED, message, details)
+
+
+# Common health checks
+
+
+async def database_health_check(pool) -> HealthStatus:
+    """
+    Health check for database connection pools.
+
+    Args:
+        pool: Database connection pool
+
+    Returns:
+        HealthStatus indicating database health
+    """
+    try:
+        # Try to acquire a connection
+        if hasattr(pool, "acquire"):
+            # asyncpg style
+            async with pool.acquire() as conn:
+                # Try a simple query
+                if hasattr(conn, "fetchval"):
+                    await conn.fetchval("SELECT 1")
+                elif hasattr(conn, "execute"):
+                    await conn.execute("SELECT 1")
+        elif hasattr(pool, "ping"):
+            # Some pools have a ping method
+            await pool.ping()
+        elif hasattr(pool, "execute"):
+            # aiosqlite style
+            await pool.execute("SELECT 1")
+
+        return HealthStatus.healthy("Database connection is healthy")
+
+    except asyncio.TimeoutError:
+        return HealthStatus.unhealthy(
+            "Database connection timed out", {"error": "timeout"}
+        )
+    except Exception as e:
+        return HealthStatus.unhealthy(
+            f"Database health check failed: {str(e)}",
+            {"error": str(e), "type": type(e).__name__},
+        )
+
+
+async def http_client_health_check(
+    client, health_endpoint: str = "/health"
+) -> HealthStatus:
+    """
+    Health check for HTTP clients.
+
+    Args:
+        client: HTTP client (aiohttp or httpx)
+        health_endpoint: Endpoint to check
+
+    Returns:
+        HealthStatus indicating HTTP client health
+    """
+    try:
+        # Handle different client types
+        if hasattr(client, "get"):
+            # aiohttp or httpx style
+            if hasattr(client, "_base_url"):
+                # httpx
+                response = await client.get(health_endpoint)
+                status_code = response.status_code
+            else:
+                # aiohttp
+                async with client.get(health_endpoint) as response:
+                    status_code = response.status
+
+            if 200 <= status_code < 300:
+                return HealthStatus.healthy(
+                    f"HTTP client is healthy (status: {status_code})"
+                )
+            else:
+                return HealthStatus.unhealthy(
+                    f"HTTP health check returned {status_code}",
+                    {"status_code": status_code},
+                )
+        else:
+            # Unknown client type, assume healthy
+            return HealthStatus.healthy(
+                "HTTP client type not recognized, assuming healthy"
+            )
+
+    except asyncio.TimeoutError:
+        return HealthStatus.unhealthy(
+            "HTTP health check timed out", {"error": "timeout"}
+        )
+    except Exception as e:
+        return HealthStatus.unhealthy(
+            f"HTTP health check failed: {str(e)}",
+            {"error": str(e), "type": type(e).__name__},
+        )
+
+
+async def cache_health_check(cache) -> HealthStatus:
+    """
+    Health check for cache clients.
+
+    Args:
+        cache: Cache client (Redis, Memcached, etc.)
+
+    Returns:
+        HealthStatus indicating cache health
+    """
+    try:
+        test_key = "_health_check_test"
+        test_value = "healthy"
+
+        # Try to set and get a value
+        if hasattr(cache, "set") and hasattr(cache, "get"):
+            await cache.set(test_key, test_value)
+            result = await cache.get(test_key)
+
+            if result == test_value or (
+                isinstance(result, bytes) and result.decode() == test_value
+            ):
+                # Clean up
+                if hasattr(cache, "delete"):
+                    await cache.delete(test_key)
+                return HealthStatus.healthy("Cache is healthy")
+            else:
+                return HealthStatus.unhealthy(
+                    "Cache health check failed: value mismatch",
+                    {"expected": test_value, "got": result},
+                )
+        elif hasattr(cache, "ping"):
+            # Redis style ping
+            await cache.ping()
+            return HealthStatus.healthy("Cache is healthy (ping successful)")
+        else:
+            # Unknown cache type
+            return HealthStatus.healthy("Cache type not recognized, assuming healthy")
+
+    except asyncio.TimeoutError:
+        return HealthStatus.unhealthy(
+            "Cache health check timed out", {"error": "timeout"}
+        )
+    except Exception as e:
+        return HealthStatus.unhealthy(
+            f"Cache health check failed: {str(e)}",
+            {"error": str(e), "type": type(e).__name__},
+        )
+
+
+async def message_queue_health_check(mq) -> HealthStatus:
+    """
+    Health check for message queue clients.
+
+    Args:
+        mq: Message queue client
+
+    Returns:
+        HealthStatus indicating message queue health
+    """
+    try:
+        # RabbitMQ (aio-pika)
+        if hasattr(mq, "channel"):
+            channel = await mq.channel()
+            await channel.close()
+            return HealthStatus.healthy(
+                "Message queue is healthy (channel test successful)"
+            )
+
+        # Kafka
+        elif hasattr(mq, "producer") and hasattr(mq, "consumer"):
+            # Custom Kafka client from factory
+            # Just check if they exist
+            return HealthStatus.healthy("Kafka clients are healthy")
+
+        # Generic check
+        elif hasattr(mq, "is_closed"):
+            if not mq.is_closed:
+                return HealthStatus.healthy("Message queue connection is open")
+            else:
+                return HealthStatus.unhealthy("Message queue connection is closed")
+
+        else:
+            # Unknown MQ type
+            return HealthStatus.healthy(
+                "Message queue type not recognized, assuming healthy"
+            )
+
+    except asyncio.TimeoutError:
+        return HealthStatus.unhealthy(
+            "Message queue health check timed out", {"error": "timeout"}
+        )
+    except Exception as e:
+        return HealthStatus.unhealthy(
+            f"Message queue health check failed: {str(e)}",
+            {"error": str(e), "type": type(e).__name__},
+        )
+
+
+def create_composite_health_check(*checks: HealthCheck) -> HealthCheck:
+    """
+    Create a composite health check from multiple checks.
+
+    All checks must pass for the resource to be considered healthy.
+
+    Args:
+        *checks: Health check functions
+
+    Returns:
+        Composite health check function
+
+    Example:
+        ```python
+        health_check = create_composite_health_check(
+            lambda r: database_health_check(r.db),
+            lambda r: cache_health_check(r.cache)
+        )
+        ```
+    """
+
+    async def composite_check(resource: Any) -> HealthStatus:
+        results = []
+
+        for check in checks:
+            try:
+                result = await check(resource)
+                if isinstance(result, bool):
+                    if not result:
+                        return HealthStatus.unhealthy("Composite check failed")
+                elif isinstance(result, HealthStatus):
+                    if not result.is_healthy:
+                        return result
+                    results.append(result)
+            except Exception as e:
+                return HealthStatus.unhealthy(
+                    f"Health check error: {str(e)}", {"error": str(e)}
+                )
+
+        # All checks passed
+        return HealthStatus.healthy("All health checks passed")
+
+    return composite_check

kailash/resources/reference.py

@@ -0,0 +1,288 @@
+"""
+Resource References - JSON-serializable references to resources.
+
+This module enables resources to be referenced in JSON APIs by providing
+a serializable reference format.
+"""
+
+import json
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class ResourceReference:
+    """
+    Reference to a resource that can be resolved by the gateway.
+
+    This class provides a JSON-serializable way to reference resources
+    in API calls, solving the problem of passing non-serializable objects.
+
+    Example:
+        ```python
+        # Create a reference to a database
+        db_ref = ResourceReference(
+            type="database",
+            config={
+                "host": "localhost",
+                "database": "myapp"
+            },
+            credentials_ref="db_credentials"
+        )
+
+        # Convert to JSON
+        json_ref = db_ref.to_json()
+
+        # Use in API call
+        response = await client.execute_workflow(
+            workflow_id="process_data",
+            inputs={"data": [1, 2, 3]},
+            resources={"db": db_ref}
+        )
+        ```
+
+    Attributes:
+        type: The type of resource (database, http_client, cache, etc.)
+        config: Configuration parameters for the resource
+        credentials_ref: Optional reference to credentials in secret manager
+        name: Optional name to reference a pre-registered resource
+    """
+
+    type: str
+    config: Dict[str, Any]
+    credentials_ref: Optional[str] = None
+    name: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert to JSON-serializable dictionary.
+
+        Returns:
+            Dictionary representation of the reference
+        """
+        data = {"type": self.type, "config": self.config}
+
+        if self.credentials_ref:
+            data["credentials_ref"] = self.credentials_ref
+
+        if self.name:
+            data["name"] = self.name
+
+        return data
+
+    def to_json(self) -> str:
+        """
+        Convert to JSON string.
+
+        Returns:
+            JSON string representation
+        """
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ResourceReference":
+        """
+        Create from dictionary.
+
+        Args:
+            data: Dictionary with reference data
+
+        Returns:
+            ResourceReference instance
+        """
+        return cls(
+            type=data["type"],
+            config=data["config"],
+            credentials_ref=data.get("credentials_ref"),
+            name=data.get("name"),
+        )
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "ResourceReference":
+        """
+        Create from JSON string.
+
+        Args:
+            json_str: JSON string representation
+
+        Returns:
+            ResourceReference instance
+        """
+        return cls.from_dict(json.loads(json_str))
+
+    @classmethod
+    def for_registered_resource(cls, name: str) -> "ResourceReference":
+        """
+        Create a reference to a pre-registered resource.
+
+        This is a shorthand for referencing resources that are already
+        registered in the ResourceRegistry.
+
+        Args:
+            name: Name of the registered resource
+
+        Returns:
+            ResourceReference instance
+
+        Example:
+            ```python
+            # Reference a pre-registered database
+            db_ref = ResourceReference.for_registered_resource("main_db")
+            ```
+        """
+        return cls(type="registered", config={}, name=name)
+
+
+def create_database_reference(
+    host: str,
+    database: str,
+    backend: str = "postgresql",
+    port: Optional[int] = None,
+    credentials_ref: Optional[str] = None,
+    **kwargs,
+) -> ResourceReference:
+    """
+    Helper to create a database resource reference.
+
+    Args:
+        host: Database host
+        database: Database name
+        backend: Database backend (postgresql, mysql, sqlite)
+        port: Database port (uses default if not specified)
+        credentials_ref: Reference to credentials in secret manager
+        **kwargs: Additional configuration
+
+    Returns:
+        ResourceReference for a database
+
+    Example:
+        ```python
+        db_ref = create_database_reference(
+            host="db.example.com",
+            database="production",
+            credentials_ref="prod_db_creds"
+        )
+        ```
+    """
+    config = {"backend": backend, "host": host, "database": database, **kwargs}
+
+    if port:
+        config["port"] = port
+
+    return ResourceReference(
+        type="database", config=config, credentials_ref=credentials_ref
+    )
+
+
+def create_http_client_reference(
+    base_url: str,
+    backend: str = "aiohttp",
+    timeout: int = 30,
+    credentials_ref: Optional[str] = None,
+    headers: Optional[Dict[str, str]] = None,
+    **kwargs,
+) -> ResourceReference:
+    """
+    Helper to create an HTTP client resource reference.
+
+    Args:
+        base_url: Base URL for the HTTP client
+        backend: HTTP client backend (aiohttp, httpx)
+        timeout: Request timeout in seconds
+        credentials_ref: Reference to credentials for auth headers
+        headers: Additional headers
+        **kwargs: Additional configuration
+
+    Returns:
+        ResourceReference for an HTTP client
+
+    Example:
+        ```python
+        api_ref = create_http_client_reference(
+            base_url="https://api.example.com",
+            timeout=60,
+            credentials_ref="api_key"
+        )
+        ```
+    """
+    config = {"backend": backend, "base_url": base_url, "timeout": timeout, **kwargs}
+
+    if headers:
+        config["headers"] = headers
+
+    return ResourceReference(
+        type="http_client", config=config, credentials_ref=credentials_ref
+    )
+
+
+def create_cache_reference(
+    backend: str = "redis",
+    host: str = "localhost",
+    port: Optional[int] = None,
+    **kwargs,
+) -> ResourceReference:
+    """
+    Helper to create a cache resource reference.
+
+    Args:
+        backend: Cache backend (redis, memcached, memory)
+        host: Cache server host
+        port: Cache server port
+        **kwargs: Additional configuration
+
+    Returns:
+        ResourceReference for a cache
+
+    Example:
+        ```python
+        cache_ref = create_cache_reference(
+            backend="redis",
+            host="cache.example.com"
+        )
+        ```
+    """
+    config = {"backend": backend, "host": host, **kwargs}
+
+    if port:
+        config["port"] = port
+
+    return ResourceReference(type="cache", config=config)
+
+
+def create_message_queue_reference(
+    backend: str,
+    host: str = "localhost",
+    port: Optional[int] = None,
+    credentials_ref: Optional[str] = None,
+    **kwargs,
+) -> ResourceReference:
+    """
+    Helper to create a message queue resource reference.
+
+    Args:
+        backend: Message queue backend (rabbitmq, kafka, redis)
+        host: Message queue host
+        port: Message queue port
+        credentials_ref: Reference to credentials
+        **kwargs: Additional configuration
+
+    Returns:
+        ResourceReference for a message queue
+
+    Example:
+        ```python
+        mq_ref = create_message_queue_reference(
+            backend="rabbitmq",
+            host="mq.example.com",
+            credentials_ref="mq_creds"
+        )
+        ```
+    """
+    config = {"backend": backend, "host": host, **kwargs}
+
+    if port:
+        config["port"] = port
+
+    return ResourceReference(
+        type="message_queue", config=config, credentials_ref=credentials_ref
+    )