executors/sql_executor.py

"""
executors/sql_executor.py
--------------------------
SQL code executor for AutoDevAgent.

Runs generated SQL queries against an in-memory SQLite instance.
Before executing the query, it sets up the schema (CREATE TABLE) and
seeds realistic dummy data (INSERT INTO) that was inferred by the
Planning Agent.

Design decisions:
  - Uses Python's built-in sqlite3 — zero cost, no external deps.
  - Everything is in-memory (:memory:) — no files created, no cleanup
    needed, and each run starts with a fresh database.
  - The SQLite connection is created INSIDE the worker thread to satisfy
    SQLite's check_same_thread requirement — connection, schema setup,
    and query execution all happen in the same thread.
  - Results are formatted as a readable ASCII table for the UI and for
    the debug agent to reason about.
  - Timeout is enforced via thread.join(timeout) — if the thread is still
    alive after the timeout, we report a timeout error.

Known limitation:
  - Dummy data is LLM-generated and realistic but not exhaustive.
    Edge cases like NULL-heavy data, empty tables, or duplicate keys
    are not guaranteed. Documented in README.

Usage:
    from executors.sql_executor import SQLExecutor
    from pipeline.state import PipelineState, Language

    executor = SQLExecutor()
    # state.sql_schema and state.generated_code must be set
    result = executor.run(state)
    print(result["execution_result"].stdout)   # ASCII table of results
    print(result["execution_result"].success)  # True / False
"""

import logging
import sqlite3
import threading
import time
from typing import Any

from config import settings
from pipeline.state import (
    ExecutionResult,
    PipelineState,
    PipelineStatus,
    SQLSchema,
)

logger = logging.getLogger(__name__)


class SQLExecutor:
    """
    Executes SQL queries against a fresh in-memory SQLite instance.

    Each run() call spawns a worker thread that creates its own
    :memory: database, loads the schema and dummy data, then runs
    the generated SQL query. The connection never crosses thread
    boundaries, satisfying SQLite thread-safety requirements.

    Attributes:
        timeout: Maximum seconds allowed for the full query run.
    """

    def __init__(self) -> None:
        """Initialise with timeout from central config."""
        self.timeout: int = settings.sql_execution_timeout

    def run(self, state: PipelineState) -> dict[str, Any]:
        """
        Set up schema, seed dummy data, and execute the SQL query.

        Args:
            state: Current PipelineState. Reads: generated_code,
                   sql_schema.

        Returns:
            Partial state dict with keys:
              - "execution_result": ExecutionResult
              - "status": PipelineStatus.EXECUTING
        """
        query = state.generated_code

        if not query or not query.strip():
            logger.warning("SQLExecutor received empty query.")
            return {
                "execution_result": ExecutionResult(
                    success=False,
                    error_msg="No SQL query was provided to execute.",
                ),
                "status": PipelineStatus.EXECUTING,
            }

        logger.info("SQLExecutor starting — query length: %d chars", len(query))

        result = self._run_in_thread(query, state.sql_schema)

        return {
            "execution_result": result,
            "status": PipelineStatus.EXECUTING,
        }

    # ---------------------------------------------------------------- #
    #  Thread runner                                                    #
    # ---------------------------------------------------------------- #

    def _run_in_thread(
        self,
        query: str,
        schema: "SQLSchema | None",
    ) -> "ExecutionResult":
        """
        Run the full SQL pipeline inside a single worker thread.

        Creating the SQLite connection inside the thread (rather than
        in the caller and passing it across) avoids the
        "SQLite objects created in a thread can only be used in that
        same thread" error that occurs with check_same_thread=True.

        Args:
            query:  SQL query string to execute.
            schema: Optional SQLSchema for table setup and dummy data.

        Returns:
            ExecutionResult populated by the worker thread.
        """
        result_holder: list = []
        start = time.perf_counter()

        def worker() -> None:
            """
            Create connection, set up schema, run query — all in one thread.
            Appends a single ExecutionResult to result_holder when done.
            """
            # ── Create fresh in-memory DB ──────────────────────── #
            try:
                conn = sqlite3.connect(":memory:")
                conn.row_factory = sqlite3.Row
            except Exception as e:
                result_holder.append(ExecutionResult(
                    success=False,
                    error_msg=f"Failed to create SQLite connection: {e}",
                ))
                return

            try:
                # ── Schema setup ───────────────────────────────── #
                schema_error = _setup_schema(conn, schema)
                if schema_error:
                    result_holder.append(ExecutionResult(
                        success=False,
                        error_msg=f"Schema setup failed: {schema_error}",
                    ))
                    return

                # ── Execute query ──────────────────────────────── #
                try:
                    cursor = conn.cursor()
                    cursor.execute(query)
                    rows    = cursor.fetchall()
                    columns = [d[0] for d in cursor.description] if cursor.description else []
                    elapsed = round(time.perf_counter() - start, 3)
                    result_holder.append(ExecutionResult(
                        success=True,
                        stdout=_format_results(rows, columns),
                        exec_time=elapsed,
                    ))
                except sqlite3.Error as e:
                    elapsed = round(time.perf_counter() - start, 3)
                    result_holder.append(ExecutionResult(
                        success=False,
                        error_msg=str(e),
                        exec_time=elapsed,
                    ))
            finally:
                conn.close()

        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        thread.join(timeout=self.timeout)

        elapsed = round(time.perf_counter() - start, 3)

        # ── Timeout path ──────────────────────────────────────── #
        if thread.is_alive():
            thread.join(timeout=1)
            msg = (
                f"SQL query timed out after {self.timeout} seconds. "
                "The query may be too complex or missing an index."
            )
            logger.warning("SQLExecutor timeout after %.3fs", elapsed)
            return ExecutionResult(
                success=False,
                error_msg=msg,
                exec_time=elapsed,
            )

        # ── No result produced ────────────────────────────────── #
        if not result_holder:
            return ExecutionResult(
                success=False,
                error_msg="SQLExecutor worker produced no result.",
                exec_time=elapsed,
            )

        exec_result = result_holder[0]
        if exec_result.success:
            logger.info(
                "SQLExecutor success in %.3fs", exec_result.exec_time
            )
        else:
            logger.info(
                "SQLExecutor failed in %.3fs: %s",
                elapsed,
                exec_result.error_msg[:120],
            )

        return exec_result


# ------------------------------------------------------------------ #
#  Module-level helpers                                                #
# ------------------------------------------------------------------ #

def _setup_schema(
    conn: sqlite3.Connection,
    schema: "SQLSchema | None",
) -> "str | None":
    """
    Execute CREATE TABLE and INSERT INTO statements on the connection.

    Called from inside the worker thread, so conn is safe to use.

    Args:
        conn:   Open SQLite connection.
        schema: SQLSchema from the Planning Agent, or None.

    Returns:
        Error message string if setup failed, else None.
    """
    if not schema:
        # No schema provided — query runs against an empty DB.
        # Fine for expressions like SELECT 1 or SELECT 42 AS answer.
        logger.debug("SQLExecutor: no schema provided, running against empty DB")
        return None

    cursor = conn.cursor()

    # Execute CREATE TABLE statements
    for stmt in schema.create_statements:
        try:
            cursor.execute(stmt)
            logger.debug("SQLExecutor: CREATE executed: %s", stmt[:80])
        except sqlite3.Error as e:
            return f"CREATE TABLE failed: {e}\nStatement: {stmt}"

    # Execute INSERT INTO statements to seed dummy data
    for stmt in schema.insert_statements:
        try:
            cursor.execute(stmt)
        except sqlite3.Error as e:
            return f"INSERT failed: {e}\nStatement: {stmt}"

    conn.commit()
    logger.info(
        "SQLExecutor: schema ready — %d table(s), %d row(s) inserted",
        len(schema.create_statements),
        len(schema.insert_statements),
    )
    return None


def _format_results(
    rows: list,
    columns: list,
) -> str:
    """
    Format SQLite query results as a readable ASCII table.

    The output is shown in the UI and passed to the debug agent so it
    can reason about whether the results are logically correct.

    Args:
        rows:    List of sqlite3.Row objects from fetchall().
        columns: List of column name strings.

    Returns:
        Multi-line string with headers, separator, and data rows.
        Returns a descriptive message for empty result sets.

    Example output:
        name  | salary
        ------|-------
        Carol | 110000.0
        Eve   | 105000.0
        Alice | 95000.0
        (3 rows)
    """
    if not columns:
        return "(Query executed successfully — no columns returned)"

    if not rows:
        return f"({', '.join(columns)})\n(0 rows)"

    # Convert each row to a list of strings for width calculation
    str_rows = [[str(row[col]) for col in columns] for row in rows]

    # Column widths = max of header length and all cell value lengths
    col_widths = [
        max(len(col), max(len(r[i]) for r in str_rows))
        for i, col in enumerate(columns)
    ]

    # Build header, separator, and data rows
    header    = " | ".join(col.ljust(col_widths[i]) for i, col in enumerate(columns))
    separator = "-+-".join("-" * w for w in col_widths)
    data_rows = [
        " | ".join(cell.ljust(col_widths[i]) for i, cell in enumerate(row))
        for row in str_rows
    ]

    count_label = f"({len(rows)} row)" if len(rows) == 1 else f"({len(rows)} rows)"

    return "\n".join([header, separator, *data_rows, count_label])