sql_tool.py

"""
sql_tool.py — a guarded SQL execution tool for the natural-language agent.

Safety design:
  1. Only SELECT statements are permitted (regex + keyword blocklist).
  2. The connection itself is opened in SQLite's URI read-only mode, so even
     a crafted statement that slipped past validation cannot write to disk.
  3. Row limit is enforced server-side to avoid dumping huge result sets.
"""

import re
import sqlite3

DB_PATH = "data/ecommerce.db"
MAX_ROWS = 200

FORBIDDEN_KEYWORDS = [
    "insert", "update", "delete", "drop", "alter", "create", "attach",
    "detach", "pragma", "vacuum", "replace", "truncate",
]

SCHEMA_DESCRIPTION = """
Tables (real Olist Brazilian e-commerce data, ~99,000 orders from 2016-2018):

customers(customer_id TEXT PK, customer_unique_id TEXT, customer_city TEXT, customer_state TEXT)
  - customer_state is a 2-letter Brazilian state code (e.g. SP, RJ, MG)

orders(order_id TEXT PK, customer_id TEXT FK -> customers, order_status TEXT,
       order_purchase_timestamp TEXT, order_delivered_customer_date TEXT,
       order_estimated_delivery_date TEXT)
  - order_status is one of: delivered, shipped, canceled, invoiced, processing, unavailable, approved, created
  - dates are TEXT in 'YYYY-MM-DD HH:MM:SS' format; some delivery dates are NULL (not yet delivered)

order_items(order_id TEXT FK -> orders, order_item_id INTEGER, product_id TEXT FK -> products,
            seller_id TEXT FK -> sellers, price REAL, freight_value REAL)

payments(order_id TEXT FK -> orders, payment_type TEXT, payment_installments INTEGER, payment_value REAL)
  - payment_type is one of: credit_card, boleto, voucher, debit_card

reviews(order_id TEXT FK -> orders, review_score INTEGER, review_comment_message TEXT, review_creation_date TEXT)
  - review_score is 1-5; review_comment_message is often NULL (many reviews have no written comment) and is in Portuguese

products(product_id TEXT PK, product_category TEXT, product_weight_g REAL)
  - product_category is in English (translated from the original Portuguese)

sellers(seller_id TEXT PK, seller_city TEXT, seller_state TEXT)

Revenue for a line item = order_items.price (freight_value is shipping cost, separate).
This is real anonymized data — expect some missing values and inconsistencies.
"""


def _is_safe_select(sql: str) -> bool:
    stripped = sql.strip().rstrip(";").strip()
    if not re.match(r"(?is)^\s*(with\b.*?\bselect|select)\b", stripped):
        return False
    lowered = stripped.lower()
    for kw in FORBIDDEN_KEYWORDS:
        if re.search(rf"\b{kw}\b", lowered):
            return False
    if ";" in stripped:  # no stacked statements
        return False
    return True


def execute_sql(query: str) -> str:
    """Execute a read-only SELECT query against the e-commerce database and return results."""
    if not _is_safe_select(query):
        return "Rejected: only single, read-only SELECT statements are permitted."

    try:
        # Open strictly read-only via URI mode — belt-and-braces beyond the regex check
        conn = sqlite3.connect(f"file:{DB_PATH}?mode=ro", uri=True)
        conn.row_factory = sqlite3.Row
        cur = conn.cursor()
        cur.execute(query)
        rows = cur.fetchmany(MAX_ROWS)
        columns = [d[0] for d in cur.description] if cur.description else []
        conn.close()

        if not rows:
            return "Query ran successfully but returned no rows."

        # Format as a simple markdown-ish table for the model to read/summarize
        header = " | ".join(columns)
        sep = " | ".join(["---"] * len(columns))
        body = "\n".join(" | ".join(str(v) for v in row) for row in rows)
        truncation_note = f"\n(showing first {MAX_ROWS} rows)" if len(rows) == MAX_ROWS else ""
        return f"{header}\n{sep}\n{body}{truncation_note}"

    except Exception as e:
        return f"SQL execution error: {e}"