Source code for operators.bedrock.safe_sql_executor

"""
SafeSQLExecutorOperator - Execute SQL with validation and limits.

This operator provides secure SQL execution with:
- SQL validation before execution (DDL/DML rejection)
- LIMIT clause enforcement
- Configurable return format (dataframe, dict, records)
- Security event logging

Usage:
    from operators.bedrock import SafeSQLExecutorOperator

    execute_query = SafeSQLExecutorOperator(
        task_id="execute_query",
        sql="SELECT * FROM my_table WHERE status = 'active'",
        database="my_database",
        max_limit=100,
    )

Author: yuan@ready.net
Date: 2026-01-29
"""

from typing import List, Optional, Tuple

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults




class SafeSQLExecutorOperator(BaseOperator):
    """
    Execute SQL query with safety validation.

    This operator validates SQL queries before execution to prevent
    dangerous operations (DDL/DML) and enforces query limits.

    :param sql: SQL query to execute (can be templated)
    :param database: Athena database name
    :param max_limit: Maximum allowed LIMIT value (default 1000)
    :param default_limit: Default LIMIT if none specified (default 100)
    :param dangerous_patterns: Custom SQL patterns to block (optional)
    :param return_format: Output format - 'dataframe', 'dict', or 'records' (default)
    :param require_limit: Whether to require a LIMIT clause (default True)
    """

    # Operator UI color (light blue for data operations)
    ui_color = "#87CEEB"

    # Templatable fields
    template_fields = ["sql", "database"]

    @apply_defaults
    def __init__(
        self,
        sql: str,
        database: str,
        max_limit: int = 1000,
        default_limit: int = 100,
        dangerous_patterns: Optional[List[Tuple[str, str]]] = None,
        return_format: str = "records",
        require_limit: bool = True,
        *args,
        **kwargs,
    ):
        super(SafeSQLExecutorOperator, self).__init__(*args, **kwargs)
        self.sql = sql
        self.database = database
        self.max_limit = max_limit
        self.default_limit = default_limit
        self.dangerous_patterns = dangerous_patterns
        self.return_format = return_format
        self.require_limit = require_limit



    def execute(self, context):
        """Execute the SQL query with validation."""
        import awswrangler as wr

        from util.bedrock.safe_sql import enforce_limit, validate_sql

        # Layer 3: Validate SQL
        is_valid, error_msg = validate_sql(
            self.sql,
            dangerous_patterns=self.dangerous_patterns,
            require_limit=self.require_limit,
        )
        if not is_valid:
            self.log.error(f"SQL validation failed: {error_msg}")
            raise ValueError(f"SQL validation failed: {error_msg}")

        # Layer 4: Enforce limit
        safe_sql = enforce_limit(
            self.sql,
            default_limit=self.default_limit,
            max_limit=self.max_limit,
        )

        self.log.info(f"Executing safe SQL: {safe_sql[:200]}...")

        # Execute query
        try:
            df = wr.athena.read_sql_query(safe_sql, database=self.database)
            self.log.info(f"Query returned {len(df)} rows")
        except Exception as e:
            self.log.error(f"Query execution failed: {str(e)}")
            raise

        # Return in requested format
        if self.return_format == "dataframe":
            return df
        elif self.return_format == "dict":
            return df.to_dict()
        else:  # records (default)
            return df.to_dict("records")