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

kailash/nodes/code/python.py

@@ -389,7 +389,9 @@ class CodeExecutor:
                 f"Error position: {' ' * (e.offset - 1) if e.offset else ''}^"
             )
 
-    def execute_code(self, code: str, inputs: dict[str, Any]) -> dict[str, Any]:
+    def execute_code(
+        self, code: str, inputs: dict[str, Any], node_instance=None
+    ) -> dict[str, Any]:
         """Execute Python code with given inputs.
 
         Args:
@@ -476,6 +478,43 @@ class CodeExecutor:
                 except ImportError:
                     logger.warning(f"Module {module_name} not available")
 
+        # Add global utility functions to namespace
+        try:
+            from kailash.utils.data_paths import (
+                get_data_path,
+                get_input_data_path,
+                get_output_data_path,
+            )
+
+            namespace["get_input_data_path"] = get_input_data_path
+            namespace["get_output_data_path"] = get_output_data_path
+            namespace["get_data_path"] = get_data_path
+        except ImportError:
+            logger.warning(
+                "Could not import data path utilities - functions will not be available in PythonCodeNode execution"
+            )
+
+        # Add workflow context functions if node instance is available
+        if node_instance and hasattr(node_instance, "get_workflow_context"):
+            # Bind the actual node methods
+            namespace["get_workflow_context"] = node_instance.get_workflow_context
+            namespace["set_workflow_context"] = node_instance.set_workflow_context
+        else:
+            # Add placeholder functions that warn about unavailability
+            def _get_workflow_context(key: str, default=None):
+                logger.warning(
+                    "get_workflow_context() is not available in PythonCodeNode execution context. Node instance not provided."
+                )
+                return default
+
+            def _set_workflow_context(key: str, value):
+                logger.warning(
+                    "set_workflow_context() is not available in PythonCodeNode execution context. Node instance not provided."
+                )
+
+            namespace["get_workflow_context"] = _get_workflow_context
+            namespace["set_workflow_context"] = _set_workflow_context
+
         # Add sanitized inputs
         namespace.update(sanitized_inputs)
 
@@ -1222,7 +1261,9 @@ class PythonCodeNode(Node):
         try:
             if self.code:
                 # Execute code string
-                outputs = self.executor.execute_code(self.code, kwargs)
+                outputs = self.executor.execute_code(
+                    self.code, kwargs, node_instance=self
+                )
                 # Return 'result' variable if it exists, otherwise all outputs
                 if "result" in outputs:
                     return {"result": outputs["result"]}
@@ -1454,7 +1495,7 @@ class PythonCodeNode(Node):
         """
         # Execute directly based on execution type
         if self.code:
-            outputs = self.executor.execute_code(self.code, inputs)
+            outputs = self.executor.execute_code(self.code, inputs, node_instance=self)
             return outputs.get("result", outputs)
         elif self.function:
             wrapper = FunctionWrapper(self.function, self.executor)

kailash/nodes/data/async_sql.py

@@ -489,6 +489,7 @@ class PostgreSQLAdapter(DatabaseAdapter):
                     or "DELETE" in query_upper
                     or "INSERT" in query_upper
                 )
+                and "SELECT" not in query_upper
                 and "RETURNING" not in query_upper
                 and fetch_mode == FetchMode.ALL
             ):
@@ -527,6 +528,7 @@ class PostgreSQLAdapter(DatabaseAdapter):
                         or "DELETE" in query_upper
                         or "INSERT" in query_upper
                     )
+                    and "SELECT" not in query_upper
                     and "RETURNING" not in query_upper
                     and fetch_mode == FetchMode.ALL
                 ):
@@ -2617,33 +2619,53 @@ class AsyncSQLDatabaseNode(AsyncNode):
 
     async def cleanup(self):
         """Clean up database connections."""
+        try:
+            # Check if we have a running event loop
+            loop = asyncio.get_running_loop()
+            if loop.is_closed():
+                # Event loop is closing, skip cleanup
+                return
+        except RuntimeError:
+            # No event loop, skip cleanup
+            return
+
         # Rollback any active transaction
         if self._active_transaction and self._adapter:
             try:
-                await self._adapter.rollback_transaction(self._active_transaction)
-            except Exception:
+                await asyncio.wait_for(
+                    self._adapter.rollback_transaction(self._active_transaction),
+                    timeout=1.0,
+                )
+            except (Exception, asyncio.TimeoutError):
                 pass  # Best effort cleanup
             self._active_transaction = None
 
         if self._adapter and self._connected:
-            if self._share_pool and self._pool_key:
-                # Decrement reference count for shared pool
-                async with self._get_pool_lock():
-                    if self._pool_key in self._shared_pools:
-                        adapter, ref_count = self._shared_pools[self._pool_key]
-                        if ref_count > 1:
-                            # Others still using the pool
-                            self._shared_pools[self._pool_key] = (
-                                adapter,
-                                ref_count - 1,
-                            )
-                        else:
-                            # Last reference, close the pool
-                            del self._shared_pools[self._pool_key]
-                            await adapter.disconnect()
-            else:
-                # Dedicated pool, close directly
-                await self._adapter.disconnect()
+            try:
+                if self._share_pool and self._pool_key:
+                    # Decrement reference count for shared pool with timeout
+                    async with await asyncio.wait_for(
+                        self._get_pool_lock(), timeout=1.0
+                    ):
+                        if self._pool_key in self._shared_pools:
+                            adapter, ref_count = self._shared_pools[self._pool_key]
+                            if ref_count > 1:
+                                # Others still using the pool
+                                self._shared_pools[self._pool_key] = (
+                                    adapter,
+                                    ref_count - 1,
+                                )
+                            else:
+                                # Last reference, close the pool
+                                del self._shared_pools[self._pool_key]
+                                await asyncio.wait_for(
+                                    adapter.disconnect(), timeout=1.0
+                                )
+                else:
+                    # Dedicated pool, close directly
+                    await asyncio.wait_for(self._adapter.disconnect(), timeout=1.0)
+            except (Exception, asyncio.TimeoutError):
+                pass  # Best effort cleanup
 
             self._connected = False
             self._adapter = None

kailash/nodes/edge/__init__.py

@@ -0,0 +1,36 @@
+"""Edge computing nodes for distributed processing and data management."""
+
+from .base import EdgeNode
+from .cloud_node import CloudNode
+from .coordination import EdgeCoordinationNode
+from .docker_node import DockerNode
+from .edge_data import EdgeDataNode
+from .edge_migration_node import EdgeMigrationNode
+from .edge_monitoring_node import EdgeMonitoringNode
+from .edge_state import EdgeStateMachine
+from .edge_warming_node import EdgeWarmingNode
+
+# Phase 4.4 Integration & Testing nodes
+from .kubernetes_node import KubernetesNode
+from .platform_node import PlatformNode
+from .resource_analyzer_node import ResourceAnalyzerNode
+from .resource_optimizer_node import ResourceOptimizerNode
+from .resource_scaler_node import ResourceScalerNode
+
+__all__ = [
+    "EdgeNode",
+    "EdgeDataNode",
+    "EdgeStateMachine",
+    "EdgeCoordinationNode",
+    "EdgeWarmingNode",
+    "EdgeMonitoringNode",
+    "EdgeMigrationNode",
+    "ResourceAnalyzerNode",
+    "ResourceScalerNode",
+    "ResourceOptimizerNode",
+    # Phase 4.4 nodes
+    "KubernetesNode",
+    "DockerNode",
+    "CloudNode",
+    "PlatformNode",
+]

kailash/nodes/edge/base.py

@@ -0,0 +1,240 @@
+"""Base edge-aware node with location awareness and compliance routing."""
+
+import asyncio
+from typing import Any, Dict, List, Optional
+
+from kailash.edge.compliance import ComplianceRouter, DataClassification
+from kailash.edge.discovery import EdgeDiscovery, EdgeSelectionStrategy
+from kailash.edge.location import EdgeLocation
+from kailash.nodes.base_async import AsyncNode
+
+
+class EdgeNode(AsyncNode):
+    """Base node with edge computing awareness.
+
+    Extends AsyncNode with:
+    - Edge location awareness
+    - Automatic edge selection
+    - Compliance-aware routing
+    - Migration capabilities
+    """
+
+    def __init__(self, **config):
+        """Initialize edge-aware node.
+
+        Args:
+            edge_strategy: Edge selection strategy (latency|cost|balanced|compliance)
+            preferred_locations: List of preferred edge location names
+            compliance_zones: List of required compliance zones (gdpr, ccpa, etc.)
+            enable_migration: Whether to enable edge migration capabilities
+            **config: Additional node configuration
+        """
+        self.edge_strategy = EdgeSelectionStrategy(
+            config.pop("edge_strategy", "balanced")
+        )
+        self.preferred_locations = config.pop("preferred_locations", [])
+        self.compliance_zones = config.pop("compliance_zones", [])
+        self.enable_migration = config.pop("enable_migration", True)
+
+        # Check for injected infrastructure (from WorkflowBuilder)
+        edge_infrastructure = config.pop("_edge_infrastructure", None)
+
+        if edge_infrastructure:
+            # Use shared infrastructure from WorkflowBuilder
+            self.edge_discovery = edge_infrastructure.get_discovery()
+            self.compliance_router = edge_infrastructure.get_compliance_router()
+            self._shared_infrastructure = edge_infrastructure
+        else:
+            # Standalone mode - create own infrastructure (backward compatibility)
+            self.edge_discovery = EdgeDiscovery()
+            self.compliance_router = ComplianceRouter()
+            self._shared_infrastructure = None
+
+        self.current_edge: Optional[EdgeLocation] = None
+
+        super().__init__(**config)
+
+    async def initialize(self):
+        """Initialize edge infrastructure."""
+        # No need to call super().initialize() as AsyncNode doesn't have it
+
+        # Start edge discovery only if not using shared infrastructure
+        if not self._shared_infrastructure:
+            await self.edge_discovery.start_discovery()
+
+        # Select initial edge
+        self.current_edge = await self._select_edge()
+
+        if not self.current_edge:
+            raise RuntimeError("No suitable edge location found")
+
+    async def _select_edge(
+        self, data: Optional[Dict[str, Any]] = None
+    ) -> Optional[EdgeLocation]:
+        """Select optimal edge location based on strategy and constraints.
+
+        Args:
+            data: Optional data for compliance classification
+
+        Returns:
+            Selected edge location or None
+        """
+        # Get all available edges
+        edges = self.edge_discovery.get_all_edges()
+
+        # Filter by preferred locations if specified
+        if self.preferred_locations:
+            edges = [e for e in edges if e.name in self.preferred_locations]
+
+        # Filter by compliance if data provided
+        if data and self.compliance_zones:
+            data_class = self.compliance_router.classify_data(data)
+            edges = [
+                e
+                for e in edges
+                if self.compliance_router.is_compliant_location(
+                    e, data_class, self.compliance_zones
+                )
+            ]
+
+        # Apply edge selection strategy
+        if not edges:
+            return None
+
+        # If we've already filtered edges, select from the filtered list
+        # based on the strategy
+        if self.edge_strategy == EdgeSelectionStrategy.LATENCY_OPTIMAL:
+            # Select edge with lowest latency
+            return min(edges, key=lambda e: e.metrics.latency_p50_ms)
+        elif self.edge_strategy == EdgeSelectionStrategy.COST_OPTIMAL:
+            # Select edge with lowest cost
+            return min(edges, key=lambda e: e.metrics.compute_cost_per_hour)
+        elif self.edge_strategy == EdgeSelectionStrategy.COMPLIANCE_FIRST:
+            # Already filtered for compliance, pick first
+            return edges[0]
+        elif self.edge_strategy == EdgeSelectionStrategy.CAPACITY_OPTIMAL:
+            # Select edge with most capacity
+            return max(
+                edges,
+                key=lambda e: e.capabilities.cpu_cores
+                * (1 - e.metrics.cpu_utilization),
+            )
+        else:  # BALANCED or others
+            # Simple balanced selection - pick edge with best combined score
+            return min(
+                edges,
+                key=lambda e: e.metrics.latency_p50_ms
+                * e.metrics.compute_cost_per_hour,
+            )
+
+    async def migrate_to_edge(
+        self, target_edge: EdgeLocation, state_data: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """Migrate this node to a different edge location.
+
+        Args:
+            target_edge: Target edge location
+            state_data: Optional state to migrate
+
+        Returns:
+            Success status
+        """
+        if not self.enable_migration:
+            return False
+
+        try:
+            # Prepare for migration
+            await self._prepare_migration(target_edge, state_data)
+
+            # Perform migration
+            old_edge = self.current_edge
+            self.current_edge = target_edge
+
+            # Cleanup old edge
+            if old_edge:
+                await self._cleanup_edge(old_edge)
+
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Migration failed: {e}")
+            return False
+
+    async def _prepare_migration(
+        self, target_edge: EdgeLocation, state_data: Optional[Dict[str, Any]]
+    ):
+        """Prepare for edge migration."""
+        # Override in subclasses for specific preparation
+        pass
+
+    async def _cleanup_edge(self, edge: EdgeLocation):
+        """Cleanup after migrating away from an edge."""
+        # Override in subclasses for specific cleanup
+        pass
+
+    async def get_edge_metrics(self) -> Dict[str, Any]:
+        """Get current edge performance metrics."""
+        if not self.current_edge:
+            return {}
+
+        return {
+            "edge_name": self.current_edge.name,
+            "edge_region": self.current_edge.region,
+            "latency_ms": self.current_edge.metrics.latency_p50_ms,
+            "cpu_usage": self.current_edge.metrics.cpu_utilization,
+            "memory_usage": self.current_edge.metrics.memory_utilization,
+            "request_count": self.current_edge.metrics.throughput_rps,
+            "error_rate": self.current_edge.metrics.error_rate,
+        }
+
+    def is_compliant_for_data(
+        self, data: Dict[str, Any], required_zones: Optional[List[str]] = None
+    ) -> bool:
+        """Check if current edge is compliant for given data.
+
+        Args:
+            data: Data to check compliance for
+            required_zones: Override compliance zones
+
+        Returns:
+            Compliance status
+        """
+        if not self.current_edge:
+            return False
+
+        zones = required_zones or self.compliance_zones
+        if not zones:
+            return True
+
+        data_class = self.compliance_router.classify_data(data)
+        return self.compliance_router.is_compliant_location(
+            self.current_edge, data_class, zones
+        )
+
+    async def ensure_compliance(
+        self, data: Dict[str, Any], required_zones: Optional[List[str]] = None
+    ) -> bool:
+        """Ensure node is at compliant edge for data.
+
+        Migrates to compliant edge if necessary.
+
+        Args:
+            data: Data requiring compliance
+            required_zones: Override compliance zones
+
+        Returns:
+            Success status
+        """
+        zones = required_zones or self.compliance_zones
+
+        # Check current compliance
+        if self.is_compliant_for_data(data, zones):
+            return True
+
+        # Find compliant edge
+        compliant_edge = await self._select_edge(data)
+        if not compliant_edge:
+            return False
+
+        # Migrate to compliant edge
+        return await self.migrate_to_edge(compliant_edge)