Mahathi4554/sql-query-debugging

1"""2Task definitions for the SQL Query Debugging Environment.3Each task defines:4  - A real-world scenario5  - A database schema (populated with seed data)6  - A broken/suboptimal SQL query for the agent to fix7  - Expected correct result rows8  - A programmatic grader9  - Progressive hints (revealed after 2 and 4 failed attempts)10"""11 12from __future__ import annotations13import sqlite314import tempfile15from dataclasses import dataclass, field16from typing import Optional17 18 19@dataclass20class Task:21    task_id: str22    difficulty: str          # easy | medium | hard | expert23    description: str24    schema_sql: str          # CREATE TABLE + INSERT statements25    broken_query: str        # What the agent receives26    solution_query: str      # Reference solution (not shown to agent)27    expected_rows: list[dict]28    hint: str                # Hint after 2 failed attempts29    hint_advanced: str       # Stronger hint after 4 failed attempts30    max_attempts: int = 531    check_optimization: bool = False32    optimization_target_cost: int = 233 34    def setup_db(self) -> str:35        """Create a fresh SQLite DB for this episode, return path."""36        tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)37        tmp.close()38        conn = sqlite3.connect(tmp.name)39        conn.executescript(self.schema_sql)40        conn.commit()41        conn.close()42        return tmp.name43 44 45# ─── Task 1: Syntax Error Fix (Easy) ─────────────────────────────────────────46 47TASK1_SCHEMA = """48CREATE TABLE employees (49    id INTEGER PRIMARY KEY,50    name TEXT NOT NULL,51    department TEXT NOT NULL,52    salary REAL NOT NULL,53    hire_date TEXT NOT NULL54);55INSERT INTO employees VALUES56  (1, 'Alice Chen',    'Engineering', 95000, '2019-03-15'),57  (2, 'Bob Martinez',  'Engineering', 88000, '2020-07-01'),58  (3, 'Carol White',   'Marketing',   72000, '2018-11-20'),59  (4, 'David Kim',     'Engineering', 102000,'2017-05-10'),60  (5, 'Eve Johnson',   'Marketing',   68000, '2021-02-28'),61  (6, 'Frank Lee',     'HR',          61000, '2022-01-15'),62  (7, 'Grace Park',    'Engineering', 91000, '2020-09-08'),63  (8, 'Henry Brown',   'HR',          58000, '2021-06-30');64"""65 66TASK1_BROKEN = """SELECT name, salary67FORM employees68WHERE department = 'Engineering'69ORDER BY salary DESC;"""70 71TASK1_SOLUTION = """SELECT name, salary72FROM employees73WHERE department = 'Engineering'74ORDER BY salary DESC;"""75 76TASK1_EXPECTED = [77    {"name": "David Kim",    "salary": 102000.0},78    {"name": "Alice Chen",   "salary": 95000.0},79    {"name": "Grace Park",   "salary": 91000.0},80    {"name": "Bob Martinez", "salary": 88000.0},81]82 83TASK1 = Task(84    task_id="task_syntax_fix",85    difficulty="easy",86    description=(87        "A data analyst needs to list all Engineering department employees "88        "by salary (highest first). The query has a typo — fix it so it runs "89        "and returns the correct result."90    ),91    schema_sql=TASK1_SCHEMA,92    broken_query=TASK1_BROKEN,93    solution_query=TASK1_SOLUTION,94    expected_rows=TASK1_EXPECTED,95    hint="Check the FROM keyword — it appears to be misspelled as FORM.",96    hint_advanced="Replace 'FORM' with 'FROM' on the second line. That's the only change needed.",97    max_attempts=5,98)99 100 101# ─── Task 2: Logic Bug Fix (Medium) ──────────────────────────────────────────102 103TASK2_SCHEMA = """104CREATE TABLE customers (105    id INTEGER PRIMARY KEY,106    name TEXT NOT NULL,107    email TEXT NOT NULL,108    region TEXT NOT NULL109);110CREATE TABLE orders (111    id INTEGER PRIMARY KEY,112    customer_id INTEGER NOT NULL,113    amount REAL NOT NULL,114    order_date TEXT NOT NULL,115    status TEXT NOT NULL116);117INSERT INTO customers VALUES118  (1, 'Acme Corp',      'acme@example.com',     'North'),119  (2, 'Beta LLC',       'beta@example.com',      'South'),120  (3, 'Gamma Inc',      'gamma@example.com',     'North'),121  (4, 'Delta Co',       'delta@example.com',     'West'),122  (5, 'Echo Ltd',       'echo@example.com',      'South');123INSERT INTO orders VALUES124  (1,  1, 750.00,  '2024-11-15', 'completed'),125  (2,  1, 200.00,  '2024-11-20', 'completed'),126  (3,  2, 1200.00, '2024-12-01', 'completed'),127  (4,  3, 450.00,  '2024-12-10', 'completed'),128  (5,  3, 600.00,  '2024-10-05', 'completed'),129  (6,  4, 850.00,  '2024-12-20', 'completed'),130  (7,  5, 300.00,  '2024-12-22', 'completed'),131  (8,  2, 550.00,  '2024-09-01', 'completed'),132  (9,  1, 900.00,  '2024-12-18', 'completed'),133  (10, 4, 125.00,  '2024-12-19', 'pending');134"""135 136TASK2_BROKEN = """SELECT DISTINCT c.name, c.email137FROM customers c138LEFT JOIN orders o ON c.id = o.customer_id139WHERE o.amount > 50140  AND o.order_date < '2024-10-01'141  AND o.status = 'completed'142ORDER BY c.name;"""143 144TASK2_SOLUTION = """SELECT DISTINCT c.name, c.email145FROM customers c146INNER JOIN orders o ON c.id = o.customer_id147WHERE o.amount > 500148  AND o.order_date >= '2024-10-01'149  AND o.status = 'completed'150ORDER BY c.name;"""151 152TASK2_EXPECTED = [153    {"name": "Acme Corp",  "email": "acme@example.com"},154    {"name": "Beta LLC",   "email": "beta@example.com"},155    {"name": "Delta Co",   "email": "delta@example.com"},156    {"name": "Gamma Inc",  "email": "gamma@example.com"},157]158 159TASK2 = Task(160    task_id="task_logic_bug",161    difficulty="medium",162    description=(163        "Find all customers who placed at least one completed order over $500 "164        "on or after 2024-10-01. The query has multiple logic bugs — "165        "the wrong JOIN type, a wrong amount threshold, and a wrong date comparison. "166        "Fix all issues to return the correct 4 customers."167    ),168    schema_sql=TASK2_SCHEMA,169    broken_query=TASK2_BROKEN,170    solution_query=TASK2_SOLUTION,171    expected_rows=TASK2_EXPECTED,172    hint="There are 3 bugs: check the JOIN type, the amount threshold value, and the direction of the date comparison.",173    hint_advanced=(174        "Fix all three: (1) LEFT JOIN → INNER JOIN, "175        "(2) amount > 50 → amount > 500, "176        "(3) order_date < '2024-10-01' → order_date >= '2024-10-01'."177    ),178    max_attempts=7,179)180 181 182# ─── Task 3: Query Optimization (Hard) ───────────────────────────────────────183 184TASK3_SCHEMA = """185CREATE TABLE products (186    id INTEGER PRIMARY KEY,187    name TEXT NOT NULL,188    category TEXT NOT NULL,189    price REAL NOT NULL190);191CREATE TABLE sales (192    id INTEGER PRIMARY KEY,193    product_id INTEGER NOT NULL,194    quantity INTEGER NOT NULL,195    sale_date TEXT NOT NULL,196    region TEXT NOT NULL197);198INSERT INTO products VALUES199  (1,  'Widget A',   'Widgets',   29.99),200  (2,  'Widget B',   'Widgets',   49.99),201  (3,  'Gadget X',   'Gadgets',   99.99),202  (4,  'Gadget Y',   'Gadgets',   149.99),203  (5,  'Doohickey',  'Other',     19.99),204  (6,  'Thingamajig','Other',     39.99),205  (7,  'Widget C',   'Widgets',   34.99),206  (8,  'Gadget Z',   'Gadgets',   199.99);207INSERT INTO sales VALUES208  (1,  1, 10, '2024-Q4', 'North'),209  (2,  1, 15, '2024-Q4', 'South'),210  (3,  2,  8, '2024-Q4', 'North'),211  (4,  2, 12, '2024-Q4', 'West'),212  (5,  3,  5, '2024-Q4', 'North'),213  (6,  3,  7, '2024-Q4', 'South'),214  (7,  4,  3, '2024-Q4', 'North'),215  (8,  4,  4, '2024-Q4', 'West'),216  (9,  5, 20, '2024-Q4', 'South'),217  (10, 6,  9, '2024-Q4', 'North'),218  (11, 7, 11, '2024-Q4', 'West'),219  (12, 8,  2, '2024-Q4', 'North'),220  (13, 1, 18, '2024-Q3', 'North'),221  (14, 3,  6, '2024-Q3', 'South'),222  (15, 5, 25, '2024-Q3', 'West');223"""224 225TASK3_BROKEN = """SELECT226    p.name,227    p.category,228    (SELECT SUM(s.quantity)229     FROM sales s230     WHERE s.product_id = p.id231       AND s.sale_date = '2024-Q4') AS total_units_sold232FROM products p233WHERE (SELECT SUM(s.quantity)234       FROM sales s235       WHERE s.product_id = p.id236         AND s.sale_date = '2024-Q4') > 10237ORDER BY total_units_sold DESC;"""238 239TASK3_SOLUTION = """SELECT240    p.name,241    p.category,242    SUM(s.quantity) AS total_units_sold243FROM products p244INNER JOIN sales s ON p.id = s.product_id245WHERE s.sale_date = '2024-Q4'246GROUP BY p.id, p.name, p.category247HAVING SUM(s.quantity) > 10248ORDER BY total_units_sold DESC;"""249 250TASK3_EXPECTED = [251    {"name": "Widget A",  "category": "Widgets", "total_units_sold": 25},252    {"name": "Widget B",  "category": "Widgets", "total_units_sold": 20},253    {"name": "Doohickey", "category": "Other",   "total_units_sold": 20},254    {"name": "Widget C",  "category": "Widgets", "total_units_sold": 11},255    {"name": "Gadget X",  "category": "Gadgets", "total_units_sold": 12},256]257 258TASK3 = Task(259    task_id="task_optimization",260    difficulty="hard",261    description=(262        "This production query lists products with more than 10 units sold in Q4-2024, "263        "but it uses correlated subqueries that make it extremely slow at scale. "264        "Rewrite it to use a JOIN + GROUP BY + HAVING approach for efficiency. "265        "The result must match the original output exactly."266    ),267    schema_sql=TASK3_SCHEMA,268    broken_query=TASK3_BROKEN,269    solution_query=TASK3_SOLUTION,270    expected_rows=TASK3_EXPECTED,271    hint="Replace the correlated subqueries with a single INNER JOIN, GROUP BY the product columns, and use HAVING to filter.",272    hint_advanced=(273        "Structure: SELECT p.name, p.category, SUM(s.quantity) AS total_units_sold "274        "FROM products p INNER JOIN sales s ON p.id = s.product_id "275        "WHERE s.sale_date = '2024-Q4' GROUP BY p.id, p.name, p.category "276        "HAVING SUM(s.quantity) > 10 ORDER BY total_units_sold DESC;"277    ),278    max_attempts=7,279    check_optimization=True,280    optimization_target_cost=2,281)282 283 284# ─── Task 4: SQL Injection Fix (Medium-Hard) ─────────────────────────────────285# NEW TASK: A query is vulnerable to SQL injection via string concatenation.286# Agent must rewrite it to use parameterized query style (safe string escaping).287# In SQLite context: agent must use proper quoting / escaping or restructure.288 289TASK4_SCHEMA = """290CREATE TABLE users (291    id INTEGER PRIMARY KEY,292    username TEXT NOT NULL UNIQUE,293    email TEXT NOT NULL,294    role TEXT NOT NULL DEFAULT 'user',295    active INTEGER NOT NULL DEFAULT 1296);297CREATE TABLE audit_log (298    id INTEGER PRIMARY KEY,299    user_id INTEGER,300    action TEXT NOT NULL,301    performed_at TEXT NOT NULL302);303INSERT INTO users VALUES304  (1, 'alice',   'alice@example.com',   'admin', 1),305  (2, 'bob',     'bob@example.com',     'user',  1),306  (3, 'charlie', 'charlie@example.com', 'user',  1),307  (4, 'diana',   'diana@example.com',   'mod',   1),308  (5, 'eve',     'eve@example.com',     'user',  0);309INSERT INTO audit_log VALUES310  (1, 1, 'login',  '2024-12-01 09:00'),311  (2, 2, 'login',  '2024-12-01 09:05'),312  (3, 1, 'update', '2024-12-01 10:00'),313  (4, 3, 'login',  '2024-12-02 08:30'),314  (5, 2, 'logout', '2024-12-02 09:00');315"""316 317# Broken: uses string concatenation — vulnerable to SQL injection.318# Also has a logic bug: should only return ACTIVE users (active = 1).319# Agent must fix both the injection pattern AND the missing active filter.320TASK4_BROKEN = """SELECT u.id, u.username, u.email, u.role,321       COUNT(a.id) AS action_count322FROM users u323LEFT JOIN audit_log a ON u.id = a.user_id324WHERE u.username = '' OR '1'='1'325GROUP BY u.id, u.username, u.email, u.role326ORDER BY u.username;"""327 328TASK4_SOLUTION = """SELECT u.id, u.username, u.email, u.role,329       COUNT(a.id) AS action_count330FROM users u331LEFT JOIN audit_log a ON u.id = a.user_id332WHERE u.active = 1333GROUP BY u.id, u.username, u.email, u.role334ORDER BY u.username;"""335 336TASK4_EXPECTED = [337    {"id": 1, "username": "alice",   "email": "alice@example.com",   "role": "admin", "action_count": 2},338    {"id": 2, "username": "bob",     "email": "bob@example.com",     "role": "user",  "action_count": 2},339    {"id": 3, "username": "charlie", "email": "charlie@example.com", "role": "user",  "action_count": 1},340    {"id": 4, "username": "diana",   "email": "diana@example.com",   "role": "mod",   "action_count": 0},341]342 343TASK4 = Task(344    task_id="task_injection_fix",345    difficulty="medium",346    description=(347        "A login audit report query contains a classic SQL injection payload in the WHERE clause "348        "('OR 1=1') that bypasses all filtering and returns every user including inactive ones. "349        "Fix the query to: (1) remove the injection payload, and (2) correctly filter to only "350        "active users (active = 1). Return all 4 active users with their audit action counts."351    ),352    schema_sql=TASK4_SCHEMA,353    broken_query=TASK4_BROKEN,354    solution_query=TASK4_SOLUTION,355    expected_rows=TASK4_EXPECTED,356    hint="The WHERE clause has a SQL injection pattern (OR '1'='1'). Replace the entire WHERE clause to filter on u.active = 1 instead.",357    hint_advanced=(358        "Replace the WHERE clause entirely: WHERE u.active = 1. "359        "Remove the username injection pattern completely. Keep the LEFT JOIN and GROUP BY as-is."360    ),361    max_attempts=6,362)363 364 365# ─── Task 5: Missing Index & Rewrite (Expert) ─────────────────────────────────366# NEW TASK: A reporting query does multiple full-table scans because it lacks367# proper use of indexed columns and uses inefficient LIKE patterns + OR chains.368# Agent must rewrite using indexed column lookups and CTEs for clarity.369 370TASK5_SCHEMA = """371CREATE TABLE transactions (372    id INTEGER PRIMARY KEY,373    account_id INTEGER NOT NULL,374    transaction_type TEXT NOT NULL,375    amount REAL NOT NULL,376    currency TEXT NOT NULL,377    transaction_date TEXT NOT NULL,378    status TEXT NOT NULL,379    merchant TEXT380);381CREATE TABLE accounts (382    id INTEGER PRIMARY KEY,383    owner_name TEXT NOT NULL,384    account_type TEXT NOT NULL,385    country TEXT NOT NULL,386    opened_date TEXT NOT NULL387);388INSERT INTO accounts VALUES389  (1, 'Sarah Connor',  'checking', 'US', '2020-01-10'),390  (2, 'John Wick',     'savings',  'US', '2019-06-15'),391  (3, 'Ellen Ripley',  'checking', 'UK', '2021-03-20'),392  (4, 'Rick Deckard',  'savings',  'US', '2018-11-05'),393  (5, 'Dana Scully',   'checking', 'CA', '2022-07-01');394INSERT INTO transactions VALUES395  (1,  1, 'debit',  250.00, 'USD', '2024-11-01', 'completed', 'Amazon'),396  (2,  1, 'credit', 3000.00,'USD', '2024-11-05', 'completed', NULL),397  (3,  2, 'debit',  1500.00,'USD', '2024-11-10', 'completed', 'Rent'),398  (4,  2, 'credit', 5000.00,'USD', '2024-11-15', 'completed', NULL),399  (5,  3, 'debit',  89.99,  'GBP', '2024-11-03', 'completed', 'Tesco'),400  (6,  3, 'debit',  450.00, 'GBP', '2024-11-20', 'pending',   'British Gas'),401  (7,  4, 'debit',  2200.00,'USD', '2024-11-12', 'completed', 'Mortgage'),402  (8,  4, 'credit', 4500.00,'USD', '2024-11-28', 'completed', NULL),403  (9,  1, 'debit',  75.50,  'USD', '2024-12-01', 'completed', 'Starbucks'),404  (10, 5, 'credit', 2800.00,'CAD', '2024-12-05', 'completed', NULL),405  (11, 5, 'debit',  120.00, 'CAD', '2024-12-08', 'failed',    'Rogers'),406  (12, 2, 'debit',  350.00, 'USD', '2024-12-10', 'completed', 'Gym');407"""408 409# Broken: uses LIKE '%completed%' (no index use), OR chain for transaction types,410# and does not aggregate correctly — double-counts due to missing DISTINCT on join.411TASK5_BROKEN = """SELECT a.owner_name,412       a.account_type,413       a.country,414       SUM(t.amount) AS total_spent,415       COUNT(t.id) AS transaction_count416FROM accounts a417LEFT JOIN transactions t ON a.id = t.account_id418WHERE t.status LIKE '%completed%'419   OR t.status LIKE '%done%'420   OR t.status LIKE '%processed%'421GROUP BY a.id422ORDER BY total_spent DESC;"""423 424TASK5_SOLUTION = """WITH completed_txns AS (425    SELECT account_id,426           SUM(CASE WHEN transaction_type = 'debit' THEN amount ELSE 0 END) AS total_spent,427           COUNT(*) AS transaction_count428    FROM transactions429    WHERE status = 'completed'430      AND transaction_type = 'debit'431    GROUP BY account_id432)433SELECT a.owner_name,434       a.account_type,435       a.country,436       COALESCE(c.total_spent, 0) AS total_spent,437       COALESCE(c.transaction_count, 0) AS transaction_count438FROM accounts a439LEFT JOIN completed_txns c ON a.id = c.account_id440ORDER BY total_spent DESC;"""441 442TASK5_EXPECTED = [443    {"owner_name": "Rick Deckard", "account_type": "savings",  "country": "US", "total_spent": 2200.0, "transaction_count": 1},444    {"owner_name": "John Wick",    "account_type": "savings",  "country": "US", "total_spent": 1850.0, "transaction_count": 2},445    {"owner_name": "Sarah Connor", "account_type": "checking", "country": "US", "total_spent": 325.5,  "transaction_count": 2},446    {"owner_name": "Dana Scully",  "account_type": "checking", "country": "CA", "total_spent": 0.0,    "transaction_count": 0},447    {"owner_name": "Ellen Ripley", "account_type": "checking", "country": "UK", "total_spent": 89.99,  "transaction_count": 1},448]449 450TASK5 = Task(451    task_id="task_reporting_rewrite",452    difficulty="expert",453    description=(454        "A financial reporting query uses LIKE '%completed%' pattern matching (preventing index use), "455        "an inefficient OR chain for status values, and incorrectly sums ALL transaction amounts "456        "instead of only debit transactions. Rewrite it using a CTE to pre-aggregate only completed "457        "debit transactions per account, then LEFT JOIN to accounts. "458        "Return all 5 accounts ordered by total_spent descending — accounts with no completed debits should show 0."459    ),460    schema_sql=TASK5_SCHEMA,461    broken_query=TASK5_BROKEN,462    solution_query=TASK5_SOLUTION,463    expected_rows=TASK5_EXPECTED,464    hint=(465        "Three issues: (1) Use status = 'completed' exact match instead of LIKE. "466        "(2) Filter to transaction_type = 'debit' only. "467        "(3) Use a CTE or subquery to pre-aggregate before joining to accounts."468    ),469    hint_advanced=(470        "Use: WITH completed_txns AS (SELECT account_id, SUM(amount) AS total_spent, COUNT(*) AS transaction_count "471        "FROM transactions WHERE status = 'completed' AND transaction_type = 'debit' GROUP BY account_id) "472        "then LEFT JOIN this CTE to accounts and COALESCE nulls to 0."473    ),474    max_attempts=8,475    check_optimization=True,476    optimization_target_cost=3,477)478 479 480# ─── Task Registry ────────────────────────────────────────────────────────────481 482TASKS: dict[str, Task] = {483    TASK1.task_id: TASK1,484    TASK2.task_id: TASK2,485    TASK3.task_id: TASK3,486    TASK4.task_id: TASK4,487    TASK5.task_id: TASK5,488}489 490TASK_LIST = [491    {492        "task_id": t.task_id,493        "difficulty": t.difficulty,494        "description": t.description,495        "max_attempts": t.max_attempts,496        "check_optimization": t.check_optimization,497        "action_schema": {498            "type": "object",499            "properties": {500                "sql_query": {501                    "type": "string",502                    "description": "The SQL query to execute as the agent's attempt",503                }504            },505            "required": ["sql_query"],506        },507    }508    for t in [TASK1, TASK2, TASK3, TASK4, TASK5]509]