butler/tools/data_query_ops.py

Rama main · réplica de lectura del repositorio original.

"""Local data query best-effort helpers (P0-A)."""

from __future__ import annotations

import json
import logging
from typing import Any

from butler.core.best_effort import safe_best_effort

logger = logging.getLogger(__name__)


def apply_statement_timeout_safe(con: Any, *, timeout_seconds: int) -> None:
    def _run() -> None:
        con.execute(f"SET statement_timeout='{timeout_seconds}s'")

    safe_best_effort(_run, label="data_query.statement_timeout", default=None)


def close_duckdb_connection_safe(con: Any) -> None:
    def _run() -> None:
        con.close()

    safe_best_effort(_run, label="data_query.close", default=None)


def execute_readonly_duckdb_query(
    *,
    query: str,
    target: str,
    cap: int,
    max_output_chars: int,
    resolve_path: Any,
    queryable_extensions: frozenset[str],
) -> str:
    import duckdb

    con = None
    try:
        con = duckdb.connect(":memory:")
        apply_statement_timeout_safe(con, timeout_seconds=30)

        p = resolve_path(target)
        if p is None:
            return json.dumps({"error": f"file not found or unsupported: {target}"})

        safe_path = str(p)
        ext = p.suffix.lower()
        if ext in (".db", ".sqlite"):
            con.execute("ATTACH ? AS src (TYPE sqlite, READ_ONLY)", [safe_path])
            con.execute("USE src")
        elif ext in (".csv", ".tsv"):
            sep = "\t" if ext == ".tsv" else ","
            con.execute(
                "CREATE TABLE data AS SELECT * FROM read_csv_auto(?, delim=?)",
                [safe_path, sep],
            )
        elif ext in (".json", ".jsonl"):
            con.execute(
                "CREATE TABLE data AS SELECT * FROM read_json_auto(?)",
                [safe_path],
            )
        elif ext == ".parquet":
            con.execute(
                "CREATE TABLE data AS SELECT * FROM read_parquet(?)",
                [safe_path],
            )
        elif ext not in queryable_extensions:
            return json.dumps({"error": f"file not found or unsupported: {target}"})

        result = con.execute(query).fetchdf()
        if len(result) > cap:
            result = result.head(cap)
            truncated = True
        else:
            truncated = False

        text = result.to_string(index=False, max_rows=cap)
        if len(text) > max_output_chars:
            text = text[:max_output_chars] + "\n…(truncated)"

        columns = list(result.columns)
        return json.dumps(
            {
                "ok": True,
                "columns": columns,
                "rows": len(result),
                "truncated": truncated,
                "text": text,
            },
            ensure_ascii=False,
        )
    except Exception as exc:
        logger.warning("data_query execution error: %s", exc)
        err_msg = str(exc)
        if len(err_msg) > 200:
            err_msg = err_msg[:200] + "…"
        return json.dumps({"error": f"query failed: {err_msg}"})
    finally:
        if con is not None:
            close_duckdb_connection_safe(con)