kailash 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

kailash/database/execution_pipeline.py

@@ -0,0 +1,499 @@
+"""Database execution pipeline for clean separation of concerns.
+
+This module provides a pipeline-based approach to database operations,
+separating permission checking, query execution, and data masking into
+clear, testable stages.
+"""
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+from kailash.access_control import NodePermission, UserContext
+from kailash.sdk_exceptions import NodeExecutionError
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExecutionContext:
+    """Context for database execution pipeline."""
+
+    query: str
+    parameters: Optional[Union[Dict[str, Any], List[Any]]] = None
+    user_context: Optional[UserContext] = None
+    node_name: str = "unknown_node"
+    result_format: str = "dict"
+    runtime_context: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class ExecutionResult:
+    """Result from database execution pipeline."""
+
+    data: Any
+    row_count: int
+    columns: List[str]
+    execution_time: float
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class PipelineStage(ABC):
+    """Abstract base class for pipeline stages."""
+
+    @abstractmethod
+    async def process(
+        self, context: ExecutionContext, result: Optional[ExecutionResult] = None
+    ) -> Optional[ExecutionResult]:
+        """Process this stage of the pipeline.
+
+        Args:
+            context: Execution context
+            result: Result from previous stage (None for first stage)
+
+        Returns:
+            Result to pass to next stage, or None to stop pipeline
+        """
+        pass
+
+    @abstractmethod
+    def get_stage_name(self) -> str:
+        """Get the name of this pipeline stage."""
+        pass
+
+
+class PermissionCheckStage(PipelineStage):
+    """Pipeline stage for checking user permissions."""
+
+    def __init__(self, access_control_manager=None):
+        """Initialize permission check stage.
+
+        Args:
+            access_control_manager: Access control manager for permission checks
+        """
+        self.access_control_manager = access_control_manager
+        self.logger = logging.getLogger(f"{__name__}.PermissionCheckStage")
+
+    async def process(
+        self, context: ExecutionContext, result: Optional[ExecutionResult] = None
+    ) -> Optional[ExecutionResult]:
+        """Check user permissions before query execution."""
+        # Skip if no access control or no user context
+        if not self.access_control_manager or not context.user_context:
+            self.logger.debug(
+                "Skipping permission check - no access control or user context"
+            )
+            return result
+
+        # Check execute permission
+        decision = self.access_control_manager.check_node_access(
+            context.user_context,
+            context.node_name,
+            NodePermission.EXECUTE,
+            context.runtime_context,
+        )
+
+        if not decision.allowed:
+            raise NodeExecutionError(f"Access denied: {decision.reason}")
+
+        self.logger.debug(
+            f"Permission granted for {context.node_name}: {decision.reason}"
+        )
+        return result
+
+    def get_stage_name(self) -> str:
+        """Get stage name."""
+        return "permission_check"
+
+
+class QueryValidationStage(PipelineStage):
+    """Pipeline stage for validating SQL queries."""
+
+    def __init__(self, validation_rules: Optional[Dict[str, Any]] = None):
+        """Initialize query validation stage.
+
+        Args:
+            validation_rules: Custom validation rules
+        """
+        self.validation_rules = validation_rules or {}
+        self.logger = logging.getLogger(f"{__name__}.QueryValidationStage")
+
+    async def process(
+        self, context: ExecutionContext, result: Optional[ExecutionResult] = None
+    ) -> Optional[ExecutionResult]:
+        """Validate query for security and safety."""
+        if not context.query:
+            raise NodeExecutionError("Query cannot be empty")
+
+        # Basic SQL injection checks
+        self._validate_query_safety(context.query)
+
+        self.logger.debug(f"Query validation passed for: {context.query[:100]}...")
+        return result
+
+    def _validate_query_safety(self, query: str) -> None:
+        """Validate query for potential security issues."""
+        if not query:
+            return
+
+        query_upper = query.upper().strip()
+
+        # Check for dangerous operations
+        dangerous_keywords = [
+            "DROP",
+            "DELETE",
+            "TRUNCATE",
+            "ALTER",
+            "CREATE",
+            "GRANT",
+            "REVOKE",
+            "EXEC",
+            "EXECUTE",
+            "SHUTDOWN",
+            "BACKUP",
+            "RESTORE",
+        ]
+
+        import re
+
+        for keyword in dangerous_keywords:
+            pattern = r"\b" + re.escape(keyword) + r"\b"
+            if re.search(pattern, query_upper):
+                self.logger.warning(
+                    f"Query contains potentially dangerous keyword: {keyword}"
+                )
+                # In production, you might want to block these entirely
+                # raise NodeExecutionError(f"Query contains forbidden keyword: {keyword}")
+
+    def get_stage_name(self) -> str:
+        """Get stage name."""
+        return "query_validation"
+
+
+class QueryExecutionStage(PipelineStage):
+    """Pipeline stage for executing SQL queries."""
+
+    def __init__(self, query_executor):
+        """Initialize query execution stage.
+
+        Args:
+            query_executor: Object that can execute queries (engine, connection, etc.)
+        """
+        self.query_executor = query_executor
+        self.logger = logging.getLogger(f"{__name__}.QueryExecutionStage")
+
+    async def process(
+        self, context: ExecutionContext, result: Optional[ExecutionResult] = None
+    ) -> Optional[ExecutionResult]:
+        """Execute the SQL query."""
+        start_time = time.time()
+
+        try:
+            # This is where the actual query execution happens
+            # The implementation depends on whether it's sync or async
+            if hasattr(self.query_executor, "execute_query"):
+                # Custom executor interface
+                query_result = await self.query_executor.execute_query(
+                    context.query, context.parameters, context.result_format
+                )
+            else:
+                # Fallback - assume it's a callable
+                query_result = await self.query_executor(
+                    context.query, context.parameters
+                )
+
+            execution_time = time.time() - start_time
+
+            # Format the result
+            if isinstance(query_result, dict):
+                # Structured result
+                return ExecutionResult(
+                    data=query_result.get("data", []),
+                    row_count=query_result.get("row_count", 0),
+                    columns=query_result.get("columns", []),
+                    execution_time=execution_time,
+                    metadata=query_result.get("metadata", {}),
+                )
+            else:
+                # Raw result - format it
+                return ExecutionResult(
+                    data=query_result,
+                    row_count=(
+                        len(query_result) if isinstance(query_result, list) else 1
+                    ),
+                    columns=[],
+                    execution_time=execution_time,
+                )
+
+        except Exception as e:
+            execution_time = time.time() - start_time
+            self.logger.error(
+                f"Query execution failed after {execution_time:.3f}s: {e}"
+            )
+            raise NodeExecutionError(f"Database query failed: {e}") from e
+
+    def get_stage_name(self) -> str:
+        """Get stage name."""
+        return "query_execution"
+
+
+class DataMaskingStage(PipelineStage):
+    """Pipeline stage for applying data masking based on user attributes."""
+
+    def __init__(self, access_control_manager=None):
+        """Initialize data masking stage.
+
+        Args:
+            access_control_manager: Access control manager with masking capabilities
+        """
+        self.access_control_manager = access_control_manager
+        self.logger = logging.getLogger(f"{__name__}.DataMaskingStage")
+
+    async def process(
+        self, context: ExecutionContext, result: Optional[ExecutionResult] = None
+    ) -> Optional[ExecutionResult]:
+        """Apply data masking based on user attributes."""
+        if not result or not result.data:
+            return result
+
+        # Skip if no access control or no user context
+        if not self.access_control_manager or not context.user_context:
+            self.logger.debug(
+                "Skipping data masking - no access control or user context"
+            )
+            return result
+
+        # Skip if not dict format (masking only works on structured data)
+        if context.result_format != "dict" or not isinstance(result.data, list):
+            self.logger.debug("Skipping data masking - data format not supported")
+            return result
+
+        # Apply masking to each row
+        masked_data = []
+        for row in result.data:
+            if isinstance(row, dict):
+                # Apply masking if access control manager supports it
+                if hasattr(self.access_control_manager, "apply_data_masking"):
+                    masked_row = self.access_control_manager.apply_data_masking(
+                        context.user_context, context.node_name, row
+                    )
+                    masked_data.append(masked_row)
+                else:
+                    masked_data.append(row)
+            else:
+                masked_data.append(row)
+
+        # Return result with masked data
+        return ExecutionResult(
+            data=masked_data,
+            row_count=result.row_count,
+            columns=result.columns,
+            execution_time=result.execution_time,
+            metadata=result.metadata,
+        )
+
+    def get_stage_name(self) -> str:
+        """Get stage name."""
+        return "data_masking"
+
+
+class DatabaseExecutionPipeline:
+    """Pipeline for executing database operations with clean separation of concerns.
+
+    This pipeline provides:
+    - Permission checking
+    - Query validation
+    - Query execution
+    - Data masking
+
+    Example:
+        >>> pipeline = DatabaseExecutionPipeline(
+        ...     access_control_manager=access_manager,
+        ...     query_executor=my_executor
+        ... )
+        >>>
+        >>> context = ExecutionContext(
+        ...     query="SELECT * FROM users",
+        ...     user_context=user,
+        ...     node_name="user_query"
+        ... )
+        >>>
+        >>> result = await pipeline.execute(context)
+    """
+
+    def __init__(
+        self,
+        access_control_manager=None,
+        query_executor=None,
+        validation_rules: Optional[Dict[str, Any]] = None,
+        custom_stages: Optional[List[PipelineStage]] = None,
+    ):
+        """Initialize database execution pipeline.
+
+        Args:
+            access_control_manager: Access control manager for permissions and masking
+            query_executor: Object that can execute database queries
+            validation_rules: Custom validation rules for queries
+            custom_stages: Additional custom pipeline stages
+        """
+        self.access_control_manager = access_control_manager
+        self.query_executor = query_executor
+        self.logger = logging.getLogger(f"{__name__}.DatabaseExecutionPipeline")
+
+        # Build pipeline stages
+        self.stages: List[PipelineStage] = []
+
+        # 1. Permission check
+        self.stages.append(PermissionCheckStage(access_control_manager))
+
+        # 2. Query validation
+        self.stages.append(QueryValidationStage(validation_rules))
+
+        # 3. Custom stages (before execution)
+        if custom_stages:
+            for stage in custom_stages:
+                if stage.get_stage_name() != "query_execution":
+                    self.stages.append(stage)
+
+        # 4. Query execution
+        if query_executor:
+            self.stages.append(QueryExecutionStage(query_executor))
+
+        # 5. Data masking
+        self.stages.append(DataMaskingStage(access_control_manager))
+
+        # 6. Custom stages (after execution)
+        if custom_stages:
+            for stage in custom_stages:
+                if stage.get_stage_name() == "post_processing":
+                    self.stages.append(stage)
+
+        self.logger.info(f"Initialized pipeline with {len(self.stages)} stages")
+
+    async def execute(self, context: ExecutionContext) -> ExecutionResult:
+        """Execute the full database pipeline.
+
+        Args:
+            context: Execution context with query, user, etc.
+
+        Returns:
+            Execution result with data, timing, etc.
+
+        Raises:
+            NodeExecutionError: If any stage fails
+        """
+        self.logger.debug(f"Starting pipeline execution for {context.node_name}")
+
+        result = None
+        pipeline_start = time.time()
+
+        try:
+            # Execute each stage in sequence
+            for i, stage in enumerate(self.stages):
+                stage_start = time.time()
+
+                try:
+                    result = await stage.process(context, result)
+                    stage_time = time.time() - stage_start
+
+                    self.logger.debug(
+                        f"Stage {i+1}/{len(self.stages)} ({stage.get_stage_name()}) "
+                        f"completed in {stage_time:.3f}s"
+                    )
+
+                    # Allow stages to stop the pipeline
+                    if result is None and stage.get_stage_name() != "permission_check":
+                        self.logger.warning(
+                            f"Pipeline stopped at stage: {stage.get_stage_name()}"
+                        )
+                        break
+
+                except Exception as e:
+                    self.logger.error(
+                        f"Pipeline failed at stage {stage.get_stage_name()}: {e}"
+                    )
+                    raise
+
+            pipeline_time = time.time() - pipeline_start
+            self.logger.info(f"Pipeline execution completed in {pipeline_time:.3f}s")
+
+            # Ensure we have a result
+            if result is None:
+                result = ExecutionResult(
+                    data=[],
+                    row_count=0,
+                    columns=[],
+                    execution_time=pipeline_time,
+                )
+
+            return result
+
+        except Exception as e:
+            pipeline_time = time.time() - pipeline_start
+            self.logger.error(
+                f"Pipeline execution failed after {pipeline_time:.3f}s: {e}"
+            )
+            raise
+
+    def add_stage(self, stage: PipelineStage, position: Optional[int] = None) -> None:
+        """Add a custom stage to the pipeline.
+
+        Args:
+            stage: Pipeline stage to add
+            position: Position to insert at (None = append)
+        """
+        if position is None:
+            self.stages.append(stage)
+        else:
+            self.stages.insert(position, stage)
+
+        self.logger.info(
+            f"Added stage {stage.get_stage_name()} at position {position or len(self.stages)}"
+        )
+
+    def remove_stage(self, stage_name: str) -> bool:
+        """Remove a stage from the pipeline.
+
+        Args:
+            stage_name: Name of stage to remove
+
+        Returns:
+            True if stage was found and removed
+        """
+        initial_count = len(self.stages)
+        self.stages = [s for s in self.stages if s.get_stage_name() != stage_name]
+        removed = len(self.stages) < initial_count
+
+        if removed:
+            self.logger.info(f"Removed stage {stage_name}")
+
+        return removed
+
+    def get_stage_info(self) -> List[Dict[str, str]]:
+        """Get information about all pipeline stages.
+
+        Returns:
+            List of stage information dictionaries
+        """
+        return [
+            {
+                "name": stage.get_stage_name(),
+                "type": type(stage).__name__,
+            }
+            for stage in self.stages
+        ]
+
+
+# Export components
+__all__ = [
+    "ExecutionContext",
+    "ExecutionResult",
+    "PipelineStage",
+    "PermissionCheckStage",
+    "QueryValidationStage",
+    "QueryExecutionStage",
+    "DataMaskingStage",
+    "DatabaseExecutionPipeline",
+]