sql_debugger_env.py

"""
SQL Query Debugger & Optimizer — OpenEnv Environment
======================================================
A real-world environment where an AI agent is given broken, inefficient,
or insecure SQL queries and must:
  - Fix syntax errors
  - Correct logical bugs (wrong JOIN, missing GROUP BY, etc.)
  - Optimize for performance (remove N+1 patterns, add proper indexes hints)
  - Detect & fix SQL injection vulnerabilities

This is a task real data analysts and backend engineers do every single day.

Tasks:
  easy   — Fix obvious syntax errors in simple SELECT queries
  medium — Fix logic bugs (wrong JOINs, incorrect aggregations, missing clauses)
  hard   — Fix + optimize + secure (injection detection, N+1 queries, subquery rewrites)

Why this beats other environments:
  - 100% deterministic graders (execute real SQL in SQLite, compare results)
  - Rich partial rewards at every step (syntax → logic → performance → security)
  - Real engineering pain point that Meta/HF engineers deal with daily
  - Novel domain — not in OpenEnv Hub yet
"""

from __future__ import annotations

import re
import sqlite3
import textwrap
from typing import Any, Dict, List, Literal, Optional, Tuple

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Database Schema (in-memory SQLite — fully reproducible)
# ---------------------------------------------------------------------------

SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS users (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    email TEXT UNIQUE NOT NULL,
    department TEXT NOT NULL,
    salary REAL NOT NULL,
    hire_date TEXT NOT NULL,
    manager_id INTEGER REFERENCES users(id)
);

CREATE TABLE IF NOT EXISTS orders (
    id INTEGER PRIMARY KEY,
    user_id INTEGER NOT NULL REFERENCES users(id),
    product TEXT NOT NULL,
    amount REAL NOT NULL,
    status TEXT NOT NULL CHECK(status IN ('pending','completed','cancelled')),
    created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS products (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    category TEXT NOT NULL,
    price REAL NOT NULL,
    stock INTEGER NOT NULL DEFAULT 0
);
"""

SEED_SQL = """
INSERT INTO users VALUES
(1,'Alice Chen','alice@co.com','Engineering',95000,'2021-03-01',NULL),
(2,'Bob Smith','bob@co.com','Engineering',82000,'2022-06-15',1),
(3,'Carol Jones','carol@co.com','Marketing',74000,'2020-01-10',NULL),
(4,'Dan Park','dan@co.com','Engineering',91000,'2019-08-22',1),
(5,'Eva Liu','eva@co.com','Marketing',68000,'2023-02-28',3),
(6,'Frank Wu','frank@co.com','HR',61000,'2021-11-05',NULL),
(7,'Grace Kim','grace@co.com','Engineering',103000,'2018-05-14',1),
(8,'Hank Patel','hank@co.com','Marketing',72000,'2022-09-30',3);

INSERT INTO products VALUES
(1,'Laptop Pro','Electronics',1299.99,45),
(2,'Wireless Mouse','Electronics',29.99,200),
(3,'Standing Desk','Furniture',549.00,30),
(4,'Monitor 4K','Electronics',699.99,60),
(5,'Ergonomic Chair','Furniture',399.00,25),
(6,'USB Hub','Electronics',49.99,150);

INSERT INTO orders VALUES
(1,1,'Laptop Pro',1299.99,'completed','2024-01-15'),
(2,2,'Wireless Mouse',29.99,'completed','2024-01-20'),
(3,1,'Monitor 4K',699.99,'completed','2024-02-01'),
(4,3,'Ergonomic Chair',399.00,'pending','2024-02-10'),
(5,4,'Standing Desk',549.00,'completed','2024-02-14'),
(6,2,'USB Hub',49.99,'cancelled','2024-02-20'),
(7,5,'Wireless Mouse',29.99,'completed','2024-03-01'),
(8,7,'Laptop Pro',1299.99,'completed','2024-03-05'),
(9,1,'USB Hub',49.99,'completed','2024-03-10'),
(10,3,'Laptop Pro',1299.99,'pending','2024-03-15'),
(11,4,'Wireless Mouse',29.99,'completed','2024-03-18'),
(12,6,'Ergonomic Chair',399.00,'completed','2024-03-20');
"""

def make_db() -> sqlite3.Connection:
    """Create fresh in-memory SQLite database."""
    conn = sqlite3.connect(":memory:")
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA_SQL)
    conn.executescript(SEED_SQL)
    conn.commit()
    return conn

def run_query(conn: sqlite3.Connection, sql: str) -> Tuple[List[Dict], Optional[str]]:
    """Execute SQL, return (rows, error). rows=[] on error."""
    try:
        cur = conn.execute(sql)
        rows = [dict(r) for r in cur.fetchall()]
        return rows, None
    except Exception as e:
        return [], str(e)

# ---------------------------------------------------------------------------
# Pydantic Models
# ---------------------------------------------------------------------------

class SQLChallenge(BaseModel):
    id: str
    description: str          # What the query SHOULD do (plain English)
    broken_sql: str           # The buggy SQL given to the agent
    expected_row_count: int   # How many rows the correct query returns
    expected_columns: List[str]  # Column names of correct output
    hint: str                 # Nudge without giving the answer
    difficulty: str
    bugs: List[str]           # Human-readable bug descriptions (shown to agent)


class Observation(BaseModel):
    challenge: SQLChallenge
    schema_info: str          # DDL so agent knows table structure
    current_step: int
    max_steps: int
    task: str
    previous_attempts: List[Dict[str, Any]] = Field(default_factory=list)
    instructions: str


class Action(BaseModel):
    challenge_id: str
    fixed_sql: str                        # The agent's corrected SQL
    explanation: Optional[str] = None     # Why the original was wrong
    detected_issues: List[str] = Field(default_factory=list)  # e.g. ["missing_group_by","sql_injection"]


class Reward(BaseModel):
    value: float
    breakdown: Dict[str, float]


# ---------------------------------------------------------------------------
# Challenges Dataset
# ---------------------------------------------------------------------------

CHALLENGES: List[Dict] = [

    # =========== EASY ===========
    {
        "id": "sq001",
        "difficulty": "easy",
        "description": "Get the names and salaries of all employees in the Engineering department, ordered by salary descending.",
        "broken_sql": """
            SELCT name, salary FORM users
            WHER department = 'Engineering'
            ORDER BY salary DESC
        """.strip(),
        "expected_row_count": 4,
        "expected_columns": ["name", "salary"],
        "hint": "Look carefully at the SQL keywords — any typos?",
        "bugs": [
            "Typo: SELCT should be SELECT",
            "Typo: FORM should be FROM",
            "Typo: WHER should be WHERE",
        ],
    },
    {
        "id": "sq002",
        "difficulty": "easy",
        "description": "Count how many orders each user has made, showing user_id and their order count.",
        "broken_sql": """
            SELECT user_id, COUNT(*) as order_count
            FROM orders
        """.strip(),
        "expected_row_count": 7,
        "expected_columns": ["user_id", "order_count"],
        "hint": "When using COUNT with a non-aggregate column, something is missing.",
        "bugs": ["Missing GROUP BY user_id — without it, SQLite returns only 1 row instead of per-user counts"],
    },
    {
        "id": "sq003",
        "difficulty": "easy",
        "description": "Get the total revenue from all completed orders.",
        "broken_sql": """
            SELECT SUM(amount) AS total_revenue
            FROM orders
            WHERE status = 'complete'
        """.strip(),
        "expected_row_count": 1,
        "expected_columns": ["total_revenue"],
        "hint": "Check the exact value used in the WHERE condition against the schema.",
        "bugs": ["Wrong status value: 'complete' should be 'completed' (as defined in the CHECK constraint)"],
    },

    # =========== MEDIUM ===========
    {
        "id": "sq004",
        "difficulty": "medium",
        "description": "Get each user's name and the total amount they've spent on completed orders. Include users who have no completed orders (show 0 for them).",
        "broken_sql": """
            SELECT u.name, SUM(o.amount) AS total_spent
            FROM orders o
            INNER JOIN users u ON u.id = o.user_id
            WHERE o.status = 'completed'
            GROUP BY u.name
        """.strip(),
        "expected_row_count": 8,
        "expected_columns": ["name", "total_spent"],
        "hint": "The problem statement says 'Include users who have no completed orders' — does INNER JOIN do that?",
        "bugs": [
            "INNER JOIN excludes users with no completed orders — should be LEFT JOIN from users to orders",
            "The WHERE clause further removes non-matching rows — it should move to an ON clause or be handled with COALESCE",
        ],
    },
    {
        "id": "sq005",
        "difficulty": "medium",
        "description": "Find all Engineering employees who earn more than the average salary of their own department.",
        "broken_sql": """
            SELECT name, salary
            FROM users
            WHERE department = 'Engineering'
            AND salary > (SELECT AVG(salary) FROM users)
        """.strip(),
        "expected_row_count": 2,
        "expected_columns": ["name", "salary"],
        "hint": "The subquery calculates something — but is it the average of the right group?",
        "bugs": [
            "Subquery uses AVG(salary) across ALL departments, not just Engineering",
            "Fix: WHERE department = 'Engineering' AND salary > (SELECT AVG(salary) FROM users WHERE department = 'Engineering')",
        ],
    },
    {
        "id": "sq006",
        "difficulty": "medium",
        "description": "Get the top 3 products by total revenue from completed orders.",
        "broken_sql": """
            SELECT product, SUM(amount) AS revenue
            FROM orders
            WHERE status = 'completed'
            GROUP BY product
            ORDER BY amount DESC
            LIMIT 3
        """.strip(),
        "expected_row_count": 3,
        "expected_columns": ["product", "revenue"],
        "hint": "The ORDER BY is using a column — but which column should you order by to rank by total revenue?",
        "bugs": [
            "ORDER BY amount DESC orders by the raw column, not the aggregated revenue",
            "Fix: ORDER BY revenue DESC (or ORDER BY SUM(amount) DESC)",
        ],
    },

    # =========== HARD ===========
    {
        "id": "sq007",
        "difficulty": "hard",
        "description": "For each department, show the department name, number of employees, average salary, and highest salary. Only include departments with more than 1 employee.",
        "broken_sql": """
            SELECT department,
                   COUNT(id),
                   AVG(salary) AS avg_salary,
                   MAX(salary)
            FROM users
            GROUP BY department
            HAVING COUNT(*) > 1
            ORDER BY avg_salary
        """.strip(),
        "expected_row_count": 2,
        "expected_columns": ["department", "COUNT(id)", "avg_salary", "MAX(salary)"],
        "hint": "The query runs but has style/clarity issues AND a subtle ordering issue. Also check column aliases.",
        "bugs": [
            "COUNT(id) and MAX(salary) have no aliases — hard to read, unpredictable column names in results",
            "ORDER BY avg_salary is ascending by default — descending is typically more useful for salary reports",
            "Minor: COUNT(id) should be COUNT(*) or COUNT(1) for clarity",
        ],
    },
    {
        "id": "sq008",
        "difficulty": "hard",
        "description": "SECURITY: A web app builds this query using user input for `dept`. Find and fix the SQL injection vulnerability. The query should return employee names for a given department.",
        "broken_sql": """
            SELECT name FROM users WHERE department = '\" + dept + \"'
        """.strip(),
        "expected_row_count": 4,
        "expected_columns": ["name"],
        "hint": "String concatenation of user input into SQL = SQL injection. The fix is parameterized queries.",
        "bugs": [
            "SQL injection vulnerability: user input `dept` is directly concatenated into the query string",
            "Fix: Use parameterized query: SELECT name FROM users WHERE department = ? with parameters=(dept,)",
            "For this environment, rewrite as: SELECT name FROM users WHERE department = 'Engineering'",
        ],
    },
    {
        "id": "sq009",
        "difficulty": "hard",
        "description": "Rewrite this inefficient N+1 style subquery pattern into a single efficient JOIN query. Get user names and their most recent order's product name and amount.",
        "broken_sql": """
            SELECT
                u.name,
                (SELECT product FROM orders WHERE user_id = u.id ORDER BY created_at DESC LIMIT 1) AS last_product,
                (SELECT amount FROM orders WHERE user_id = u.id ORDER BY created_at DESC LIMIT 1) AS last_amount
            FROM users u
        """.strip(),
        "expected_row_count": 8,
        "expected_columns": ["name", "last_product", "last_amount"],
        "hint": "Two correlated subqueries hit the orders table twice per user. Use a CTE or subquery with ROW_NUMBER() or a ranked join instead.",
        "bugs": [
            "N+1 query pattern: two correlated subqueries each scan orders once per user = O(n) full scans",
            "Fix: Use a single subquery that gets the latest order per user, then LEFT JOIN it",
        ],
    },
]

CHALLENGE_MAP: Dict[str, Dict] = {c["id"]: c for c in CHALLENGES}

# Correct reference SQL for each challenge (used internally for grading)
REFERENCE_SQL: Dict[str, str] = {
    "sq001": "SELECT name, salary FROM users WHERE department = 'Engineering' ORDER BY salary DESC",
    "sq002": "SELECT user_id, COUNT(*) as order_count FROM orders GROUP BY user_id",
    "sq003": "SELECT SUM(amount) AS total_revenue FROM orders WHERE status = 'completed'",
    "sq004": """
        SELECT u.name, COALESCE(SUM(o.amount), 0) AS total_spent
        FROM users u
        LEFT JOIN orders o ON u.id = o.user_id AND o.status = 'completed'
        GROUP BY u.id, u.name
    """,
    "sq005": """
        SELECT name, salary FROM users
        WHERE department = 'Engineering'
        AND salary > (SELECT AVG(salary) FROM users WHERE department = 'Engineering')
    """,
    "sq006": """
        SELECT product, SUM(amount) AS revenue
        FROM orders WHERE status = 'completed'
        GROUP BY product ORDER BY revenue DESC LIMIT 3
    """,
    "sq007": """
        SELECT department, COUNT(*) AS employee_count, AVG(salary) AS avg_salary, MAX(salary) AS max_salary
        FROM users GROUP BY department HAVING COUNT(*) > 1 ORDER BY avg_salary DESC
    """,
    "sq008": "SELECT name FROM users WHERE department = 'Engineering'",
    "sq009": """
        SELECT u.name, latest.product AS last_product, latest.amount AS last_amount
        FROM users u
        LEFT JOIN (
            SELECT user_id, product, amount,
                   ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY created_at DESC) AS rn
            FROM orders
        ) latest ON latest.user_id = u.id AND latest.rn = 1
    """,
}

# =========== Task Config ===========
TASK_CONFIGS = {
    "easy":   {"challenge_ids": ["sq001", "sq002", "sq003"], "max_steps": 3},
    "medium": {"challenge_ids": ["sq004", "sq005", "sq006"], "max_steps": 3},
    "hard":   {"challenge_ids": ["sq007", "sq008", "sq009"], "max_steps": 3},
}

SCHEMA_INFO = textwrap.dedent("""
    TABLE users(id, name, email, department, salary, hire_date, manager_id)
      - department: 'Engineering' | 'Marketing' | 'HR'
      - salary: REAL
      - manager_id: FK → users.id (nullable)

    TABLE orders(id, user_id, product, amount, status, created_at)
      - status: 'pending' | 'completed' | 'cancelled'
      - user_id: FK → users.id

    TABLE products(id, name, category, price, stock)
      - category: 'Electronics' | 'Furniture'
""").strip()

INSTRUCTIONS = {
    "easy": (
        "EASY TASK — Fix SQL Syntax Errors\n"
        "Each challenge gives you a broken SQL query with typos or missing keywords.\n"
        "Fix the SQL so it runs correctly and returns the expected results.\n"
        "Set fixed_sql to your corrected query."
    ),
    "medium": (
        "MEDIUM TASK — Fix Logic Bugs\n"
        "Each query runs without error but produces WRONG results due to logic bugs:\n"
        "wrong JOIN type, wrong column in ORDER BY, missing GROUP BY, wrong subquery scope.\n"
        "Fix the SQL to return the correct, expected result set."
    ),
    "hard": (
        "HARD TASK — Fix, Optimize & Secure\n"
        "Each query has logic bugs AND performance/security issues:\n"
        "SQL injection vulnerabilities, N+1 correlated subqueries, missing aliases.\n"
        "Fix ALL issues. For SQL injection: rewrite using safe parameterized form.\n"
        "For N+1 queries: rewrite as a single efficient JOIN. List all detected_issues."
    ),
}


# ---------------------------------------------------------------------------
# Environment
# ---------------------------------------------------------------------------

class SQLDebuggerEnv:
    """OpenEnv-compliant SQL Debugger & Optimizer environment."""

    def __init__(self, task: str = "easy"):
        if task not in TASK_CONFIGS:
            raise ValueError(f"Unknown task '{task}'. Choose: {list(TASK_CONFIGS)}")
        self.task = task
        self._cfg = TASK_CONFIGS[task]
        self._challenges: List[SQLChallenge] = []
        self._step = 0
        self._done = False
        self._history: List[Dict] = []
        self._results: Dict[str, Dict] = {}
        self._db: Optional[sqlite3.Connection] = None
        self.reset()

    # ------------------------------------------------------------------
    # OpenEnv interface
    # ------------------------------------------------------------------

    def reset(self) -> Observation:
        self._db = make_db()
        self._challenges = [
            SQLChallenge(**CHALLENGE_MAP[cid])
            for cid in self._cfg["challenge_ids"]
        ]
        self._step = 0
        self._done = False
        self._history = []
        self._results = {}
        return self._make_observation(0)

    def step(self, action: Action) -> Tuple[Observation, float, bool, Dict]:
        if self._done:
            raise RuntimeError("Episode done. Call reset().")

        self._step += 1
        challenge = self._find_challenge(action.challenge_id)
        idx = self._challenge_index(action.challenge_id)

        if challenge is None or idx is None:
            reward = -0.1
            info = {"error": "invalid challenge_id", "breakdown": {}}
        else:
            breakdown = self._grade(challenge, action)
            reward = max(0.0, min(1.0, sum(breakdown.values())))
            info = {"breakdown": breakdown, "step": self._step}
            self._results[action.challenge_id] = {
                "action": action.model_dump(),
                "reward": reward,
                "breakdown": breakdown,
            }

        self._history.append({
            "step": self._step,
            "challenge_id": action.challenge_id,
            "reward": reward,
            "info": info,
        })

        next_idx = min(idx + 1 if idx is not None else 0, len(self._challenges) - 1)
        self._done = (
            self._step >= self._cfg["max_steps"]
            or len(self._results) >= len(self._challenges)
        )

        obs = self._make_observation(next_idx if not self._done else idx)
        return obs, round(reward, 4), self._done, info

    def state(self) -> Dict[str, Any]:
        return {
            "task": self.task,
            "step": self._step,
            "done": self._done,
            "challenges_total": len(self._challenges),
            "challenges_solved": len(self._results),
            "history": self._history,
        }

    def close(self):
        if self._db:
            self._db.close()

    def episode_score(self) -> float:
        if not self._results:
            return 0.0
        total = sum(r["reward"] for r in self._results.values())
        return round(total / len(self._challenges), 4)

    # ------------------------------------------------------------------
    # Grading (deterministic — runs real SQL)
    # ------------------------------------------------------------------

    def _grade(self, challenge: SQLChallenge, action: Action) -> Dict[str, float]:
        bd: Dict[str, float] = {}
        fixed_sql = action.fixed_sql.strip()

        # 1. Syntax check — does the fixed SQL run at all?
        rows, error = run_query(self._db, fixed_sql)
        if error:
            bd["syntax_error"] = 0.0
            bd["parse_penalty"] = -0.05
            return bd
        bd["syntax_ok"] = 0.20

        # 2. Row count match
        ref_rows, _ = run_query(self._db, REFERENCE_SQL[challenge.id])
        if len(rows) == len(ref_rows):
            bd["row_count_correct"] = 0.25
        elif abs(len(rows) - len(ref_rows)) <= 1:
            bd["row_count_close"] = 0.10

        # 3. Column names match
        if rows and ref_rows:
            pred_cols = set(rows[0].keys())
            ref_cols = set(ref_rows[0].keys())
            col_overlap = len(pred_cols & ref_cols) / max(len(ref_cols), 1)
            bd["columns"] = round(col_overlap * 0.15, 4)

        # 4. Data correctness — compare sorted string representations
        if rows and ref_rows and len(rows) == len(ref_rows):
            pred_str = sorted(str(sorted(r.items())) for r in rows)
            ref_str  = sorted(str(sorted(r.items())) for r in ref_rows)
            if pred_str == ref_str:
                bd["data_exact"] = 0.30
            else:
                # Partial: at least first-column values match
                pred_first = sorted(str(list(r.values())[0]) for r in rows)
                ref_first  = sorted(str(list(r.values())[0]) for r in ref_rows)
                if pred_first == ref_first:
                    bd["data_partial"] = 0.15

        # 5. Security grading (hard task: injection)
        if challenge.id == "sq008":
            sql_lower = fixed_sql.lower()
            injection_pattern = "dept"
            if injection_pattern not in fixed_sql and "concat" not in sql_lower:
                bd["injection_removed"] = 0.10
            # Reward parameterized hint in explanation
            if action.explanation and any(
                kw in action.explanation.lower()
                for kw in ["parameterized", "prepared", "placeholder", "?", "injection"]
            ):
                bd["security_explanation"] = 0.05

        # 6. Optimization grading (hard task: N+1)
        if challenge.id == "sq009":
            sql_lower = fixed_sql.lower()
            # Penalize correlated subquery pattern still present
            if "select" in sql_lower[sql_lower.find("select")+6:]:  # nested SELECT
                subq_count = sql_lower.count("select")
                if subq_count <= 2:  # CTE or single subquery = OK
                    bd["optimization_ok"] = 0.05
                else:
                    bd["n_plus_1_penalty"] = -0.10
            if "join" in sql_lower:
                bd["uses_join"] = 0.05

        # 7. Explanation quality bonus
        if action.explanation and len(action.explanation.split()) >= 8:
            bd["explanation_bonus"] = 0.05

        # 8. Bug detection quality (hard task)
        if self.task == "hard" and action.detected_issues:
            bd["issues_detected"] = min(0.05, len(action.detected_issues) * 0.02)

        return bd

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _find_challenge(self, cid: str) -> Optional[SQLChallenge]:
        for c in self._challenges:
            if c.id == cid:
                return c
        return None

    def _challenge_index(self, cid: str) -> Optional[int]:
        for i, c in enumerate(self._challenges):
            if c.id == cid:
                return i
        return None

    def _make_observation(self, idx: int) -> Observation:
        challenge = self._challenges[min(idx, len(self._challenges) - 1)]
        return Observation(
            challenge=challenge,
            schema_info=SCHEMA_INFO,
            current_step=self._step,
            max_steps=self._cfg["max_steps"],
            task=self.task,
            previous_attempts=[
                h for h in self._history
                if h["challenge_id"] == challenge.id
            ],
            instructions=INSTRUCTIONS[self.task],
        )