debug_env/server/tools/execute_sql/tool.py

"""
execute_sql tool — run a SQL query against the task's SQLite database.

Diagnostic (and repair) tool with pass_rate=0.0: the agent can use this to
inspect poisoned rows, update statuses, and clean up corrupted state.

Dangerous keywords (DROP, ATTACH, DETACH) are blocked to prevent accidental
destruction of the benchmark environment.
"""

import logging
import re
import sqlite3
from pathlib import Path

from debug_env.server.schemas import ToolResult

logger = logging.getLogger(__name__)

_BLOCKED_KEYWORDS = frozenset({"DROP", "ATTACH", "DETACH"})
_BLOCKED_RE = re.compile(
    r"\b(" + "|".join(_BLOCKED_KEYWORDS) + r")\b", re.IGNORECASE
)


def execute_sql(workdir: str, query: str) -> ToolResult:
    """
    Execute *query* against ``{workdir}/app.db`` and return the results.

    SELECT queries display rows as a plain-text table.  DML statements
    (INSERT/UPDATE/DELETE) display the affected row count.

    Args:
        workdir: Absolute path to the task working directory.
        query:   SQL query to execute.

    Returns:
        :class:`ToolResult` with ``pass_rate=0.0`` and query output in ``logs``.
    """
    if not query or not query.strip():
        return ToolResult(pass_rate=0.0, logs="No query provided.", success=False)

    # Block dangerous statements
    match = _BLOCKED_RE.search(query)
    if match:
        keyword = match.group(0).upper()
        return ToolResult(
            pass_rate=0.0,
            logs=f"SQL keyword '{keyword}' is not permitted.",
            success=False,
        )

    db_path = Path(workdir) / "app.db"
    if not db_path.exists():
        return ToolResult(
            pass_rate=0.0,
            logs="No app.db found in the workdir. This task may not use a database.",
            success=False,
        )

    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    try:
        cursor = conn.execute(query)
        if cursor.description:
            rows = cursor.fetchall()
            if not rows:
                result = "(no rows returned)"
            else:
                col_names = [d[0] for d in cursor.description]
                # Compute column widths
                widths = [len(c) for c in col_names]
                str_rows = []
                for row in rows:
                    str_row = [str(v) if v is not None else "NULL" for v in row]
                    str_rows.append(str_row)
                    for i, cell in enumerate(str_row):
                        widths[i] = max(widths[i], len(cell))

                def fmt_row(cells: list) -> str:
                    return " | ".join(cell.ljust(widths[i]) for i, cell in enumerate(cells))

                header = fmt_row(col_names)
                divider = "-+-".join("-" * w for w in widths)
                body = "\n".join(fmt_row(r) for r in str_rows)
                result = f"{header}\n{divider}\n{body}\n({len(rows)} row(s))"
        else:
            conn.commit()
            result = f"OK — {cursor.rowcount} row(s) affected."

    except sqlite3.Error as exc:
        return ToolResult(pass_rate=0.0, logs=f"SQL error: {exc}", success=False)
    finally:
        conn.close()

    logger.info("execute_sql: query=%r → %d chars output", query[:80], len(result))
    return ToolResult(pass_rate=0.0, logs=result, success=True)