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

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."""
@@ -581,7 +789,7 @@ async def quick_health_check(service_name: str) -> bool:
 @dataclass
 class HealthSummary:
     """Health summary for all checks."""
-    
+
     total_checks: int
     healthy_checks: int
     degraded_checks: int
@@ -592,7 +800,7 @@ class HealthSummary:
 
 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
@@ -600,77 +808,77 @@ class HealthCheckManager:
         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)
-        
-        # Check for status changes
-        await self._check_status_change(check_name, 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) 
+                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
-            ])
+
+            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)  
+        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
@@ -680,27 +888,29 @@ class HealthCheckManager:
             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
+            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]:
+
+    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:
@@ -712,39 +922,57 @@ class HealthCheckManager:
                         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):
+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()
-    if hasattr(manager, 'register_health_check'):
-        manager.register_health_check(health_check, interval)
-    else:
-        manager.register_check(name, health_check)
+    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):
+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."""
-    # This would need MemoryHealthCheck implementation
-    pass
+    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):
+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."""
-    # This would need CustomHealthCheck implementation  
-    pass
-
-
-def get_health_manager():
-    """Get the global health manager instance."""
-    # Return a default HealthMonitor for now
-    return HealthMonitor()
+    health_check = CustomHealthCheck(name, check_func, timeout=timeout)
+    manager = get_health_manager()
+    manager.register_health_check(health_check, interval)