kailash 0.8.4__py3-none-any.whl → 0.8.6__py3-none-any.whl

kailash/core/resilience/circuit_breaker.py

@@ -396,6 +396,16 @@ class ConnectionCircuitBreaker(Generic[T]):
         if listener in self._listeners:
             self._listeners.remove(listener)
 
+    @property
+    def success_count(self) -> int:
+        """Get number of successful calls."""
+        return self.metrics.successful_calls
+
+    @property
+    def failure_count(self) -> int:
+        """Get number of failed calls."""
+        return self.metrics.failed_calls
+
     def get_status(self) -> Dict[str, Any]:
         """Get current circuit breaker status."""
         return {
@@ -435,10 +445,16 @@ class ConnectionCircuitBreaker(Generic[T]):
 class CircuitBreakerManager:
     """Manages multiple circuit breakers for different resources."""
 
-    def __init__(self):
+    def __init__(self, performance_monitor=None):
         """Initialize circuit breaker manager."""
         self._breakers: Dict[str, ConnectionCircuitBreaker] = {}
         self._default_config = CircuitBreakerConfig()
+        self._performance_monitor = performance_monitor
+        self._patterns = {
+            "database": CircuitBreakerConfig(failure_threshold=5, recovery_timeout=60),
+            "api": CircuitBreakerConfig(failure_threshold=3, recovery_timeout=30),
+            "cache": CircuitBreakerConfig(failure_threshold=2, recovery_timeout=15),
+        }
 
     def get_or_create(
         self, name: str, config: Optional[CircuitBreakerConfig] = None
@@ -450,6 +466,63 @@ class CircuitBreakerManager:
             )
         return self._breakers[name]
 
+    def create_circuit_breaker(
+        self,
+        name: str,
+        config: Optional[CircuitBreakerConfig] = None,
+        pattern: Optional[str] = None,
+    ) -> ConnectionCircuitBreaker:
+        """Create a new circuit breaker with optional pattern-based configuration."""
+        if pattern and pattern in self._patterns:
+            config = config or self._patterns[pattern]
+        return self.get_or_create(name, config)
+
+    async def execute_with_circuit_breaker(
+        self, name: str, func: Callable, fallback: Optional[Callable] = None
+    ):
+        """Execute a function with circuit breaker protection."""
+        cb = self.get_or_create(name)
+        try:
+            result = await cb.call(func)
+            return result
+        except CircuitBreakerError:
+            if fallback:
+                if asyncio.iscoroutinefunction(fallback):
+                    return await fallback()
+                else:
+                    return fallback()
+            raise
+
+    def get_circuit_breaker(self, name: str) -> Optional[ConnectionCircuitBreaker]:
+        """Get an existing circuit breaker by name."""
+        return self._breakers.get(name)
+
+    def get_all_circuit_states(self) -> Dict[str, Dict[str, Any]]:
+        """Get the state of all circuit breakers."""
+        return {name: cb.get_status() for name, cb in self._breakers.items()}
+
+    def force_open_circuit_breaker(self, name: str) -> bool:
+        """Manually open a circuit breaker."""
+        cb = self._breakers.get(name)
+        if cb:
+            asyncio.create_task(cb.force_open("Manual override"))
+            return True
+        return False
+
+    def reset_circuit_breaker(self, name: str) -> bool:
+        """Reset a circuit breaker to closed state."""
+        cb = self._breakers.get(name)
+        if cb:
+            asyncio.create_task(cb.reset())
+            return True
+        return False
+
+    def add_global_callback(self, callback: Callable):
+        """Add a global callback for circuit breaker state changes."""
+        # For now, add to all existing breakers
+        for cb in self._breakers.values():
+            cb.add_listener(callback)
+
     def get_all_status(self) -> Dict[str, Dict[str, Any]]:
         """Get status of all circuit breakers."""
         return {name: breaker.get_status() for name, breaker in self._breakers.items()}

kailash/core/resilience/health_monitor.py

@@ -64,10 +64,33 @@ class HealthCheckResult:
     error_message: Optional[str] = None
     is_healthy: bool = field(init=False)
 
+    # Additional attributes for compatibility
+    check_name: str = field(default="", init=False)
+    message: str = field(default="", init=False)
+    error: Optional[str] = field(default=None, init=False)
+    metadata: Dict[str, Any] = field(default_factory=dict, init=False)
+
     def __post_init__(self):
-        """Calculate health status."""
+        """Calculate health status and initialize compatibility fields."""
         self.is_healthy = self.status in [HealthStatus.HEALTHY, HealthStatus.DEGRADED]
 
+        # Initialize compatibility fields
+        self.check_name = self.service_name
+        self.error = self.error_message
+        self.metadata = self.details.copy()
+
+        # Set message based on status
+        if self.status == HealthStatus.HEALTHY:
+            self.message = "Service is healthy"
+        elif self.status == HealthStatus.DEGRADED:
+            self.message = "Service is degraded but functional"
+        elif self.status == HealthStatus.UNHEALTHY:
+            self.message = (
+                f"Service is unhealthy: {self.error_message or 'Unknown error'}"
+            )
+        else:
+            self.message = "Service status unknown"
+
 
 @dataclass
 class HealthMetrics:
@@ -123,10 +146,20 @@ class HealthCheck(ABC):
 class DatabaseHealthCheck(HealthCheck):
     """Health check for database connections."""
 
-    def __init__(self, name: str, connection_string: str, **kwargs):
+    def __init__(self, name: str, database_node_or_connection_string, **kwargs):
         """Initialize database health check."""
         super().__init__(name, **kwargs)
-        self.connection_string = connection_string
+        self.check_name = name  # Required by HealthCheckManager
+
+        # Handle both database node objects and connection strings
+        if hasattr(database_node_or_connection_string, "execute"):
+            # It's a database node object
+            self.database_node = database_node_or_connection_string
+            self.connection_string = None
+        else:
+            # It's a connection string
+            self.connection_string = database_node_or_connection_string
+            self.database_node = None
 
     async def check_health(self) -> HealthCheckResult:
         """Check database health."""
@@ -134,41 +167,72 @@ class DatabaseHealthCheck(HealthCheck):
         check_id = str(uuid4())
 
         try:
-            # Import SQL node for health checking
-            from src.kailash.nodes.data.sql import SQLDatabaseNode
-
-            sql_node = SQLDatabaseNode(connection_string=self.connection_string)
+            if self.database_node:
+                # Use database node object directly
+                result = await self.database_node.execute(
+                    "SELECT 1 as health_check", "dict"
+                )
 
-            # Execute simple health check query
-            result = await asyncio.wait_for(
-                asyncio.to_thread(sql_node.execute, query="SELECT 1 as health_check"),
-                timeout=self.timeout,
-            )
+                response_time = (time.time() - start_time) * 1000
+
+                if result and result.get("success"):
+                    return HealthCheckResult(
+                        check_id=check_id,
+                        service_name=self.name,
+                        status=HealthStatus.HEALTHY,
+                        response_time_ms=response_time,
+                        details={
+                            "query_executed": True,
+                            "query_result": result.get("data", []),
+                        },
+                    )
+                else:
+                    return HealthCheckResult(
+                        check_id=check_id,
+                        service_name=self.name,
+                        status=HealthStatus.DEGRADED,
+                        response_time_ms=response_time,
+                        details={"query_executed": True, "query_result": []},
+                        error_message="Query returned no success result",
+                    )
+            else:
+                # Use connection string with SQL node
+                from src.kailash.nodes.data.sql import SQLDatabaseNode
 
-            response_time = (time.time() - start_time) * 1000
+                sql_node = SQLDatabaseNode(connection_string=self.connection_string)
 
-            if "data" in result and len(result["data"]) > 0:
-                return HealthCheckResult(
-                    check_id=check_id,
-                    service_name=self.name,
-                    status=HealthStatus.HEALTHY,
-                    response_time_ms=response_time,
-                    details={
-                        "query_executed": True,
-                        "rows_returned": len(result["data"]),
-                        "execution_time": result.get("execution_time", 0),
-                    },
-                )
-            else:
-                return HealthCheckResult(
-                    check_id=check_id,
-                    service_name=self.name,
-                    status=HealthStatus.DEGRADED,
-                    response_time_ms=response_time,
-                    details={"query_executed": True, "rows_returned": 0},
-                    error_message="Query returned no data",
+                # Execute simple health check query
+                result = await asyncio.wait_for(
+                    asyncio.to_thread(
+                        sql_node.execute, query="SELECT 1 as health_check"
+                    ),
+                    timeout=self.timeout,
                 )
 
+                response_time = (time.time() - start_time) * 1000
+
+                if "data" in result and len(result["data"]) > 0:
+                    return HealthCheckResult(
+                        check_id=check_id,
+                        service_name=self.name,
+                        status=HealthStatus.HEALTHY,
+                        response_time_ms=response_time,
+                        details={
+                            "query_executed": True,
+                            "rows_returned": len(result["data"]),
+                            "execution_time": result.get("execution_time", 0),
+                        },
+                    )
+                else:
+                    return HealthCheckResult(
+                        check_id=check_id,
+                        service_name=self.name,
+                        status=HealthStatus.DEGRADED,
+                        response_time_ms=response_time,
+                        details={"query_executed": True, "rows_returned": 0},
+                        error_message="Query returned no data",
+                    )
+
         except asyncio.TimeoutError:
             response_time = (time.time() - start_time) * 1000
             return HealthCheckResult(
@@ -196,6 +260,7 @@ class RedisHealthCheck(HealthCheck):
         """Initialize Redis health check."""
         super().__init__(name, **kwargs)
         self.redis_config = redis_config
+        self.check_name = name  # Required by HealthCheckManager
 
     async def check_health(self) -> HealthCheckResult:
         """Check Redis health."""
@@ -249,6 +314,148 @@ class RedisHealthCheck(HealthCheck):
             )
 
 
+class MemoryHealthCheck(HealthCheck):
+    """Health check for system memory usage."""
+
+    def __init__(
+        self,
+        name: str,
+        warning_threshold: float = 80.0,
+        critical_threshold: float = 95.0,
+        **kwargs,
+    ):
+        """Initialize memory health check."""
+        super().__init__(name, **kwargs)
+        self.warning_threshold = warning_threshold
+        self.critical_threshold = critical_threshold
+        self.check_name = name  # Required by HealthCheckManager
+
+    async def check_health(self) -> HealthCheckResult:
+        """Check system memory health."""
+        start_time = time.time()
+        check_id = str(uuid4())
+
+        try:
+            import psutil
+
+            memory = psutil.virtual_memory()
+            response_time = (time.time() - start_time) * 1000
+
+            # Determine status based on memory usage
+            if memory.percent >= self.critical_threshold:
+                status = HealthStatus.UNHEALTHY
+                message = f"Critical memory usage: {memory.percent:.1f}%"
+            elif memory.percent >= self.warning_threshold:
+                status = HealthStatus.DEGRADED
+                message = f"High memory usage: {memory.percent:.1f}%"
+            else:
+                status = HealthStatus.HEALTHY
+                message = f"Memory usage normal: {memory.percent:.1f}%"
+
+            return HealthCheckResult(
+                check_id=check_id,
+                service_name=self.name,
+                status=status,
+                response_time_ms=response_time,
+                details={
+                    "memory_percent": memory.percent,
+                    "total_memory": memory.total,
+                    "available_memory": memory.available,
+                    "used_memory": memory.used,
+                },
+                error_message=message if status != HealthStatus.HEALTHY else None,
+            )
+
+        except Exception as e:
+            response_time = (time.time() - start_time) * 1000
+            return HealthCheckResult(
+                check_id=check_id,
+                service_name=self.name,
+                status=HealthStatus.UNHEALTHY,
+                response_time_ms=response_time,
+                error_message=str(e),
+            )
+
+
+class CustomHealthCheck(HealthCheck):
+    """Custom health check that executes user-defined check function."""
+
+    def __init__(self, name: str, check_function: Callable, **kwargs):
+        """Initialize custom health check."""
+        super().__init__(name, **kwargs)
+        self.check_function = check_function
+        self.check_name = name  # Required by HealthCheckManager
+
+    async def check_health(self) -> HealthCheckResult:
+        """Execute custom health check function."""
+        start_time = time.time()
+        check_id = str(uuid4())
+
+        try:
+            # Execute the custom check function
+            if asyncio.iscoroutinefunction(self.check_function):
+                result = await asyncio.wait_for(
+                    self.check_function(), timeout=self.timeout
+                )
+            else:
+                result = await asyncio.wait_for(
+                    asyncio.to_thread(self.check_function), timeout=self.timeout
+                )
+
+            response_time = (time.time() - start_time) * 1000
+
+            # Handle different result formats
+            if isinstance(result, bool):
+                status = HealthStatus.HEALTHY if result else HealthStatus.UNHEALTHY
+                message = "Check passed" if result else "Check failed"
+                details = {"result": result}
+            elif isinstance(result, dict):
+                # Expect dict with status, message, metadata
+                status_str = result.get("status", "healthy").lower()
+                if status_str == "healthy":
+                    status = HealthStatus.HEALTHY
+                elif status_str == "degraded":
+                    status = HealthStatus.DEGRADED
+                else:
+                    status = HealthStatus.UNHEALTHY
+
+                message = result.get("message", "Custom check completed")
+                details = result.get("metadata", {})
+            else:
+                # Assume success if we get any non-false result
+                status = HealthStatus.HEALTHY
+                message = "Custom check completed"
+                details = {"result": str(result)}
+
+            return HealthCheckResult(
+                check_id=check_id,
+                service_name=self.name,
+                status=status,
+                response_time_ms=response_time,
+                details=details,
+                error_message=None if status == HealthStatus.HEALTHY else message,
+            )
+
+        except asyncio.TimeoutError:
+            response_time = (time.time() - start_time) * 1000
+            return HealthCheckResult(
+                check_id=check_id,
+                service_name=self.name,
+                status=HealthStatus.UNHEALTHY,
+                response_time_ms=response_time,
+                error_message=f"Custom health check timeout after {self.timeout}s",
+            )
+        except Exception as e:
+            response_time = (time.time() - start_time) * 1000
+            return HealthCheckResult(
+                check_id=check_id,
+                service_name=self.name,
+                status=HealthStatus.UNHEALTHY,
+                response_time_ms=response_time,
+                error_message=str(e),
+            )
+
+
 class HTTPHealthCheck(HealthCheck):
     """Health check for HTTP endpoints."""
 
@@ -257,6 +464,7 @@ class HTTPHealthCheck(HealthCheck):
         super().__init__(name, **kwargs)
         self.url = url
         self.expected_status = expected_status
+        self.check_name = name  # Required by HealthCheckManager
 
     async def check_health(self) -> HealthCheckResult:
         """Check HTTP endpoint health."""
@@ -576,3 +784,195 @@ async def quick_health_check(service_name: str) -> bool:
         return result.is_healthy if result else False
     except Exception:
         return False
+
+
+@dataclass
+class HealthSummary:
+    """Health summary for all checks."""
+
+    total_checks: int
+    healthy_checks: int
+    degraded_checks: int
+    unhealthy_checks: int
+    overall_status: HealthStatus
+    details: List[HealthCheckResult]
+
+
+class HealthCheckManager:
+    """Manager for orchestrating multiple health checks with configuration."""
+
+    def __init__(self, config: Dict[str, Any]):
+        """Initialize health check manager with configuration."""
+        self.config = config
+        self.enabled = config.get("enabled", True)
+        self.default_interval = config.get("default_interval", 30.0)
+        self.parallel_checks = config.get("parallel_checks", True)
+        self.max_concurrent_checks = config.get("max_concurrent_checks", 10)
+
+        self.health_checks: Dict[str, HealthCheck] = {}
+        self.check_intervals: Dict[str, float] = {}
+        self.last_results: Dict[str, HealthCheckResult] = {}
+        self.history: Dict[str, List[HealthCheckResult]] = {}
+        self.status_change_callbacks: List[Callable] = []
+        self._running = False
+
+    def register_health_check(self, health_check: HealthCheck, interval: float = None):
+        """Register a health check with optional interval."""
+        check_name = health_check.check_name
+        self.health_checks[check_name] = health_check
+        self.check_intervals[check_name] = interval or self.default_interval
+        self.history[check_name] = []
+
+    async def run_health_check(self, check_name: str) -> HealthCheckResult:
+        """Run a specific health check."""
+        if check_name not in self.health_checks:
+            raise ValueError(f"Health check '{check_name}' not found")
+
+        health_check = self.health_checks[check_name]
+        result = await health_check.check_health()
+
+        # Check for status changes before storing new result
+        await self._check_status_change(check_name, result)
+
+        # Store result
+        self.last_results[check_name] = result
+        self.history[check_name].append(result)
+
+        return result
+
+    async def run_all_health_checks(self) -> List[HealthCheckResult]:
+        """Run all registered health checks."""
+        if not self.health_checks:
+            return []
+
+        if self.parallel_checks:
+            # Run checks in parallel
+            tasks = [
+                self.run_health_check(check_name)
+                for check_name in self.health_checks.keys()
+            ]
+
+            # Limit concurrency
+            semaphore = asyncio.Semaphore(self.max_concurrent_checks)
+
+            async def run_with_semaphore(task):
+                async with semaphore:
+                    return await task
+
+            results = await asyncio.gather(
+                *[run_with_semaphore(task) for task in tasks]
+            )
+        else:
+            # Run checks sequentially
+            results = []
+            for check_name in self.health_checks.keys():
+                result = await self.run_health_check(check_name)
+                results.append(result)
+
+        return results
+
+    async def get_health_summary(self) -> HealthSummary:
+        """Get summary of all health checks."""
+        results = await self.run_all_health_checks()
+
+        healthy_count = sum(1 for r in results if r.status == HealthStatus.HEALTHY)
+        degraded_count = sum(1 for r in results if r.status == HealthStatus.DEGRADED)
+        unhealthy_count = sum(1 for r in results if r.status == HealthStatus.UNHEALTHY)
+
+        # Determine overall status
+        if unhealthy_count > 0:
+            overall_status = HealthStatus.UNHEALTHY
+        elif degraded_count > 0:
+            overall_status = HealthStatus.DEGRADED
+        elif healthy_count > 0:
+            overall_status = HealthStatus.HEALTHY
+        else:
+            overall_status = HealthStatus.UNKNOWN
+
+        return HealthSummary(
+            total_checks=len(results),
+            healthy_checks=healthy_count,
+            degraded_checks=degraded_count,
+            unhealthy_checks=unhealthy_count,
+            overall_status=overall_status,
+            details=results,
+        )
+
+    def add_status_change_callback(self, callback: Callable):
+        """Add callback for status changes."""
+        self.status_change_callbacks.append(callback)
+
+    def get_health_history(
+        self, check_name: str, limit: int = None
+    ) -> List[HealthCheckResult]:
+        """Get health check history for a specific check."""
+        history = self.history.get(check_name, [])
+        if limit:
+            return history[-limit:]
+        return history
+
+    async def _check_status_change(self, check_name: str, result: HealthCheckResult):
+        """Check if status has changed and notify callbacks."""
+        if check_name in self.last_results:
+            previous = self.last_results[check_name]
+            if previous.status != result.status:
+                # Status changed, notify callbacks
+                for callback in self.status_change_callbacks:
+                    try:
+                        await callback(check_name, result)
+                    except Exception as e:
+                        logger.error(f"Error in status change callback: {e}")
+
+    async def shutdown(self):
+        """Shutdown the health check manager."""
+        self._running = False
+        # Any cleanup logic here
+
+
+# Global health manager instance for convenience functions
+_global_health_manager: Optional[HealthCheckManager] = None
+
+
+def get_health_manager() -> HealthCheckManager:
+    """Get the global health manager instance."""
+    global _global_health_manager
+    if _global_health_manager is None:
+        config = {
+            "enabled": True,
+            "default_interval": 30.0,
+            "parallel_checks": True,
+            "max_concurrent_checks": 10,
+        }
+        _global_health_manager = HealthCheckManager(config)
+    return _global_health_manager
+
+
+# Add convenience functions for registering health checks
+async def register_database_health_check(
+    name: str, database_node, interval: float = 30.0
+):
+    """Register a database health check with global manager."""
+    health_check = DatabaseHealthCheck(name, database_node)
+    manager = get_health_manager()
+    manager.register_health_check(health_check, interval)
+
+
+async def register_memory_health_check(
+    name: str,
+    warning_threshold: float = 80.0,
+    critical_threshold: float = 95.0,
+    interval: float = 30.0,
+):
+    """Register a memory health check with global manager."""
+    health_check = MemoryHealthCheck(name, warning_threshold, critical_threshold)
+    manager = get_health_manager()
+    manager.register_health_check(health_check, interval)
+
+
+async def register_custom_health_check(
+    name: str, check_func: Callable, interval: float = 30.0, timeout: float = 10.0
+):
+    """Register a custom health check with global manager."""
+    health_check = CustomHealthCheck(name, check_func, timeout=timeout)
+    manager = get_health_manager()
+    manager.register_health_check(health_check, interval)

kailash/edge/compliance.py

@@ -785,6 +785,39 @@ class ComplianceRouter:
         """Get recent compliance decisions from audit log."""
         return self.audit_log[-limit:]
 
+    def is_compliant_location(
+        self,
+        location: "EdgeLocation",
+        data_class: DataClassification,
+        required_zones: List[str],
+    ) -> bool:
+        """Check if a location is compliant for given data class and zones.
+
+        Args:
+            location: Edge location to check
+            data_class: Classification of the data
+            required_zones: Required compliance zones
+
+        Returns:
+            True if location is compliant
+        """
+        # Avoid circular import
+        from kailash.edge.location import EdgeRegion
+
+        # Check if location has all required compliance zones
+        location_zones = [z.value for z in location.compliance_zones]
+
+        # For GDPR compliance, PII/EU_PERSONAL data must be in EU regions or GDPR-compliant zones
+        if "gdpr" in required_zones and data_class in [
+            DataClassification.PII,
+            DataClassification.EU_PERSONAL,
+        ]:
+            # Check if location has GDPR compliance zone
+            return "gdpr" in location_zones
+
+        # For other cases, check if location has the required zones
+        return all(zone in location_zones for zone in required_zones)
+
     def get_compliance_summary(self) -> Dict[str, Any]:
         """Get summary of compliance decisions and performance."""
         if not self.audit_log: