kailash 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

kailash/nodes/code/python.py

@@ -50,6 +50,7 @@ import ast
 import importlib.util
 import inspect
 import logging
+import os
 import resource
 import traceback
 from collections.abc import Callable
@@ -94,6 +95,7 @@ ALLOWED_MODULES = {
     "matplotlib",
     "seaborn",
     "plotly",
+    "array",  # Required by numpy internally
     # File processing modules
     "csv",  # For CSV file processing
     "mimetypes",  # For MIME type detection
@@ -419,12 +421,59 @@ class CodeExecutor:
         }
 
         # Add allowed modules
-        for module_name in self.allowed_modules:
-            try:
-                module = importlib.import_module(module_name)
-                namespace[module_name] = module
-            except ImportError:
-                logger.warning(f"Module {module_name} not available")
+        # Check if we're running under coverage to avoid instrumentation conflicts
+        import sys
+
+        if "coverage" in sys.modules:
+            # Under coverage, use lazy loading for problematic modules
+            problematic_modules = {
+                "numpy",
+                "scipy",
+                "sklearn",
+                "pandas",
+                "matplotlib",
+                "seaborn",
+                "plotly",
+                "array",
+            }
+            safe_modules = self.allowed_modules - problematic_modules
+
+            # Eagerly load safe modules
+            for module_name in safe_modules:
+                try:
+                    module = importlib.import_module(module_name)
+                    namespace[module_name] = module
+                except ImportError:
+                    logger.warning(f"Module {module_name} not available")
+
+            # Add lazy loader for problematic modules
+            class LazyModuleLoader:
+                def __getattr__(self, name):
+                    if name in problematic_modules:
+                        return importlib.import_module(name)
+                    raise AttributeError(f"Module {name} not found")
+
+            # Make problematic modules available through lazy loading
+            for module_name in problematic_modules:
+                try:
+                    # Try to import the module directly
+                    module = importlib.import_module(module_name)
+                    namespace[module_name] = module
+                except ImportError:
+                    # If import fails, use lazy loader as fallback
+                    namespace[module_name] = LazyModuleLoader()
+        else:
+            # Normal operation - eagerly load all modules
+            for module_name in self.allowed_modules:
+                try:
+                    # Skip scipy in CI due to version conflicts
+                    if module_name == "scipy" and os.environ.get("CI"):
+                        logger.warning("Skipping scipy import in CI environment")
+                        continue
+                    module = importlib.import_module(module_name)
+                    namespace[module_name] = module
+                except ImportError:
+                    logger.warning(f"Module {module_name} not available")
 
         # Add sanitized inputs
         namespace.update(sanitized_inputs)

kailash/nodes/data/async_sql.py

@@ -431,6 +431,7 @@ class PostgreSQLAdapter(DatabaseAdapter):
         fetch_mode: FetchMode = FetchMode.ALL,
         fetch_size: Optional[int] = None,
         transaction: Optional[Any] = None,
+        parameter_types: Optional[dict[str, str]] = None,
     ) -> Any:
         """Execute query and return results."""
         # Convert dict params to positional for asyncpg
@@ -440,8 +441,10 @@ class PostgreSQLAdapter(DatabaseAdapter):
             import json
 
             query_params = []
+            param_names = []  # Track parameter names for type mapping
             for i, (key, value) in enumerate(params.items(), 1):
                 query = query.replace(f":{key}", f"${i}")
+                param_names.append(key)
                 # For PostgreSQL, lists should remain as lists for array operations
                 # Only convert dicts to JSON strings
                 if isinstance(value, dict):
@@ -449,6 +452,24 @@ class PostgreSQLAdapter(DatabaseAdapter):
                 query_params.append(value)
             params = query_params
 
+            # Apply parameter type casts if provided
+            if parameter_types:
+                # Build a query with explicit type casts
+                for i, param_name in enumerate(param_names, 1):
+                    if param_name in parameter_types:
+                        pg_type = parameter_types[param_name]
+                        # Replace $N with $N::type in the query
+                        query = query.replace(f"${i}", f"${i}::{pg_type}")
+
+        else:
+            # For positional params, apply type casts if provided
+            if parameter_types and isinstance(params, (list, tuple)):
+                # Build query with type casts for positional parameters
+                for i, param_type in parameter_types.items():
+                    if isinstance(i, int) and 0 <= i < len(params):
+                        # Replace $N with $N::type
+                        query = query.replace(f"${i+1}", f"${i+1}::{param_type}")
+
         # Ensure params is a list/tuple for asyncpg
         if params is None:
             params = []
@@ -1270,6 +1291,13 @@ class AsyncSQLDatabaseNode(AsyncNode):
                 default=False,
                 description="Whether to allow administrative SQL commands (CREATE, DROP, etc.)",
             ),
+            NodeParameter(
+                name="parameter_types",
+                type=dict,
+                required=False,
+                default=None,
+                description="Optional PostgreSQL type hints for parameters (e.g., {'role_id': 'text', 'metadata': 'jsonb'})",
+            ),
             NodeParameter(
                 name="retry_config",
                 type=Any,
@@ -1532,6 +1560,9 @@ class AsyncSQLDatabaseNode(AsyncNode):
                 "result_format", self.config.get("result_format", "dict")
             )
             user_context = inputs.get("user_context")
+            parameter_types = inputs.get(
+                "parameter_types", self.config.get("parameter_types")
+            )
 
             if not query:
                 raise NodeExecutionError("No query provided")
@@ -1576,8 +1607,12 @@ class AsyncSQLDatabaseNode(AsyncNode):
                 fetch_mode=fetch_mode,
                 fetch_size=fetch_size,
                 user_context=user_context,
+                parameter_types=parameter_types,
             )
 
+            # Ensure all data is JSON-serializable (safety net for adapter inconsistencies)
+            result = self._ensure_serializable(result)
+
             # Format results based on requested format
             formatted_data = self._format_results(result, result_format)
 
@@ -1795,6 +1830,7 @@ class AsyncSQLDatabaseNode(AsyncNode):
         fetch_mode: FetchMode,
         fetch_size: Optional[int],
         user_context: Any = None,
+        parameter_types: Optional[dict[str, str]] = None,
     ) -> Any:
         """Execute query with retry logic for transient failures.
 
@@ -1823,6 +1859,7 @@ class AsyncSQLDatabaseNode(AsyncNode):
                     params=params,
                     fetch_mode=fetch_mode,
                     fetch_size=fetch_size,
+                    parameter_types=parameter_types,
                 )
 
                 # Apply data masking if access control is enabled
@@ -2010,6 +2047,7 @@ class AsyncSQLDatabaseNode(AsyncNode):
         params: Any,
         fetch_mode: FetchMode,
         fetch_size: Optional[int],
+        parameter_types: Optional[dict[str, str]] = None,
     ) -> Any:
         """Execute query with automatic transaction management.
 
@@ -2034,6 +2072,7 @@ class AsyncSQLDatabaseNode(AsyncNode):
                 fetch_mode=fetch_mode,
                 fetch_size=fetch_size,
                 transaction=self._active_transaction,
+                parameter_types=parameter_types,
             )
         elif self._transaction_mode == "auto":
             # Auto-transaction mode
@@ -2045,6 +2084,7 @@ class AsyncSQLDatabaseNode(AsyncNode):
                     fetch_mode=fetch_mode,
                     fetch_size=fetch_size,
                     transaction=transaction,
+                    parameter_types=parameter_types,
                 )
                 await adapter.commit_transaction(transaction)
                 return result
@@ -2058,6 +2098,7 @@ class AsyncSQLDatabaseNode(AsyncNode):
                 params=params,
                 fetch_mode=fetch_mode,
                 fetch_size=fetch_size,
+                parameter_types=parameter_types,
             )
 
     @classmethod
@@ -2460,6 +2501,55 @@ class AsyncSQLDatabaseNode(AsyncNode):
 
         return modified_query, param_dict
 
+    def _ensure_serializable(self, data: Any) -> Any:
+        """Ensure all data types are JSON-serializable.
+
+        This is a safety net for cases where adapter _convert_row might not be called
+        or might miss certain data types. It recursively processes the data structure
+        to ensure datetime objects and other non-JSON-serializable types are converted.
+
+        Args:
+            data: Raw data from database adapter
+
+        Returns:
+            JSON-serializable data structure
+        """
+        if data is None:
+            return None
+        elif isinstance(data, bool):
+            return data
+        elif isinstance(data, (int, float, str)):
+            return data
+        elif isinstance(data, datetime):
+            return data.isoformat()
+        elif isinstance(data, date):
+            return data.isoformat()
+        elif hasattr(data, "total_seconds"):  # timedelta
+            return data.total_seconds()
+        elif isinstance(data, Decimal):
+            return float(data)
+        elif isinstance(data, bytes):
+            import base64
+
+            return base64.b64encode(data).decode("utf-8")
+        elif hasattr(data, "__str__") and hasattr(data, "hex"):  # UUID-like objects
+            return str(data)
+        elif isinstance(data, dict):
+            return {
+                key: self._ensure_serializable(value) for key, value in data.items()
+            }
+        elif isinstance(data, (list, tuple)):
+            return [self._ensure_serializable(item) for item in data]
+        else:
+            # For any other type, try to convert to string as fallback
+            try:
+                # Test if it's already JSON serializable
+                json.dumps(data)
+                return data
+            except (TypeError, ValueError):
+                # Not serializable, convert to string
+                return str(data)
+
     def _format_results(self, data: list[dict], result_format: str) -> Any:
         """Format query results according to specified format.