backend/app/agents/execution.py

"""
Execution Agent — Safely executes validated SQL against the tenant database.
Includes EXPLAIN-based cost estimation to prevent expensive queries.
Tracks execution time and handles errors gracefully.
"""

import time
import structlog

from app.agents.state import AgentState

logger = structlog.get_logger()

# Maximum rows a query is allowed to examine before being blocked
MAX_ROWS_EXAMINED = 100_000


def _estimate_query_cost(db_pool, sql: str) -> dict:
    """
    Run EXPLAIN on the query to estimate cost before execution.
    Returns estimated row count, whether a full table scan is detected,
    and whether the query is safe to execute.
    """
    try:
        explain_sql = f"EXPLAIN {sql.rstrip(';')}"
        rows = db_pool.execute_query(explain_sql)

        total_rows_examined = 0
        has_full_scan = False

        for row in rows:
            row_estimate = int(row.get("rows", 0) or 0)
            total_rows_examined += row_estimate
            scan_type = str(row.get("type", "")).upper()
            if scan_type == "ALL" and row_estimate > 10000:
                has_full_scan = True

        return {
            "estimated_rows": total_rows_examined,
            "has_full_scan": has_full_scan,
            "safe": total_rows_examined < MAX_ROWS_EXAMINED and not has_full_scan,
        }
    except Exception as e:
        logger.warning("explain_failed", error=str(e))
        # If EXPLAIN fails, allow the query (fail open)
        return {"estimated_rows": 0, "has_full_scan": False, "safe": True}


def execution_node(state: AgentState, db_pool) -> dict:
    """
    Execute the validated SQL query against the database.
    Runs EXPLAIN first to estimate cost and block expensive queries.
    Measures execution time and returns structured results.
    """
    sql = state.get("sanitized_sql", "") or state.get("generated_sql", "")
    trace_id = state.get("trace_id", "unknown")

    logger.info("agent_started", agent="execution", trace_id=trace_id)

    if not sql or not sql.strip():
        return {
            "query_results": [],
            "execution_time_ms": 0,
            "row_count": 0,
            "column_names": [],
            "error": "No SQL to execute",
            "error_agent": "execution",
        }

    # ── Cost Estimation (conditional to avoid doubling DB round-trips) ──
    # Only run EXPLAIN for queries without WHERE filters or complex queries,
    # since simple filtered queries are unlikely to cause full table scans.
    sql_upper = sql.upper()
    needs_cost_check = "WHERE" not in sql_upper or state.get("complexity") == "complex"

    if needs_cost_check:
        cost = _estimate_query_cost(db_pool, sql)
        if not cost["safe"]:
            logger.warning(
                "query_too_expensive",
                trace_id=trace_id,
                estimated_rows=cost["estimated_rows"],
                has_full_scan=cost["has_full_scan"],
            )
            return {
                "query_results": [],
                "execution_time_ms": 0,
                "row_count": 0,
                "column_names": [],
                "error": f"Query blocked: estimated to scan ~{cost['estimated_rows']:,} rows. "
                         f"Add WHERE filters or LIMIT to reduce scope.",
                "error_agent": "execution",
                "friendly_message": (
                    "⚠️ **Query too expensive**: This query would scan a very large number of rows. "
                    "Please add filters (WHERE clause) or a smaller LIMIT to reduce the scope."
                ),
            }
    else:
        cost = {"estimated_rows": 0, "has_full_scan": False, "safe": True}

    # ── Execute Query ────────────────────────────────────
    try:
        start_time = time.perf_counter()
        results = db_pool.execute_query(sql)
        end_time = time.perf_counter()

        execution_time_ms = round((end_time - start_time) * 1000, 2)

        # Extract column names from results
        column_names = list(results[0].keys()) if results else []

        logger.info(
            "query_executed",
            execution_time_ms=execution_time_ms,
            row_count=len(results),
            columns=len(column_names),
            estimated_rows=cost["estimated_rows"],
        )

        return {
            "query_results": results,
            "execution_time_ms": execution_time_ms,
            "row_count": len(results),
            "column_names": column_names,
        }

    except Exception as e:
        error_msg = str(e)
        logger.error("query_execution_failed", error=error_msg, sql_preview=sql[:100])

        return {
            "query_results": [],
            "execution_time_ms": 0,
            "row_count": 0,
            "column_names": [],
            "error": f"Database error: {error_msg}",
            "error_agent": "execution",
            "friendly_message": f"The query had a syntax error: {error_msg}",
        }