dkAmulet/sql-query-optimizer
0
1"""2Task definitions and deterministic graders — v2 (4 tasks).3 4Tasks:5 1. select_star_removal — easy Remove SELECT *6 2. subquery_to_join — medium IN subquery → JOIN7 3. aggregation_optimization — hard Correlated subqueries → GROUP BY8 4. cte_refactoring — expert Deeply nested subqueries → CTEs9 10All graders return (score: float, breakdown: dict, feedback: str).11Scores are always in [0.0, 1.0].12"""13from __future__ import annotations14 15import re16import sqlite317from typing import Dict, List, Tuple18 19from models import ExecutionMetrics20 21# ── Schema DDL (shown to agent) ───────────────────────────────────────────────22 23SCHEMA_DDL = """\24-- Users registered on the platform25CREATE TABLE users (26 user_id INTEGER PRIMARY KEY,27 username TEXT NOT NULL,28 email TEXT NOT NULL,29 first_name TEXT,30 last_name TEXT,31 city TEXT,32 country TEXT,33 is_active INTEGER DEFAULT 1,34 created_at TEXT35);36 37-- Two-level product taxonomy38CREATE TABLE categories (39 category_id INTEGER PRIMARY KEY,40 name TEXT NOT NULL,41 parent_category_id INTEGER42);43 44-- Product suppliers45CREATE TABLE suppliers (46 supplier_id INTEGER PRIMARY KEY,47 name TEXT NOT NULL,48 country TEXT,49 rating REAL50);51 52-- Products for sale53CREATE TABLE products (54 product_id INTEGER PRIMARY KEY,55 name TEXT NOT NULL,56 category_id INTEGER REFERENCES categories(category_id),57 supplier_id INTEGER REFERENCES suppliers(supplier_id),58 price REAL,59 sku TEXT UNIQUE,60 is_available INTEGER DEFAULT 161);62 63-- Stock levels per product64CREATE TABLE inventory (65 inventory_id INTEGER PRIMARY KEY,66 product_id INTEGER UNIQUE REFERENCES products(product_id),67 stock_qty INTEGER DEFAULT 0,68 reorder_level INTEGER DEFAULT 10,69 last_updated TEXT70);71 72-- Customer orders73CREATE TABLE orders (74 order_id INTEGER PRIMARY KEY,75 user_id INTEGER REFERENCES users(user_id),76 status TEXT, -- pending|processing|shipped|delivered|cancelled77 total_amount REAL,78 coupon_id INTEGER,79 created_at TEXT80);81 82-- Line items within an order83CREATE TABLE order_items (84 item_id INTEGER PRIMARY KEY,85 order_id INTEGER REFERENCES orders(order_id),86 product_id INTEGER REFERENCES products(product_id),87 quantity INTEGER,88 unit_price REAL89);90 91-- Product reviews92CREATE TABLE reviews (93 review_id INTEGER PRIMARY KEY,94 product_id INTEGER REFERENCES products(product_id),95 user_id INTEGER REFERENCES users(user_id),96 rating INTEGER CHECK(rating BETWEEN 1 AND 5),97 review_text TEXT,98 created_at TEXT99);100 101-- Discount coupons102CREATE TABLE coupons (103 coupon_id INTEGER PRIMARY KEY,104 code TEXT UNIQUE,105 discount_pct REAL,106 is_active INTEGER DEFAULT 1107);108 109-- Key indexes (exploit these for performance)110CREATE INDEX idx_users_active ON users(is_active);111CREATE INDEX idx_users_country ON users(country, is_active);112CREATE INDEX idx_orders_user ON orders(user_id);113CREATE INDEX idx_orders_status ON orders(status);114CREATE INDEX idx_orders_created ON orders(created_at);115CREATE INDEX idx_items_order ON order_items(order_id);116CREATE INDEX idx_items_product ON order_items(product_id);117CREATE INDEX idx_products_cat ON products(category_id);118CREATE INDEX idx_products_avail ON products(is_available);119CREATE INDEX idx_reviews_product ON reviews(product_id);120CREATE INDEX idx_inventory_prod ON inventory(product_id);121"""122 123# ── Task registry ─────────────────────────────────────────────────────────────124 125TASK_ORDER: List[str] = [126 "select_star_removal",127 "subquery_to_join",128 "aggregation_optimization",129 "cte_refactoring",130]131 132TASKS: Dict[str, dict] = {133 "select_star_removal": {134 "name": "SELECT * Elimination",135 "difficulty": "easy",136 "max_steps": 5,137 "description": (138 "The slow query uses SELECT * to fetch every column from the users table. "139 "Rewrite it to return ONLY the three required columns: user_id, username, email. "140 "The query MUST still return ALL active users (is_active = 1)."141 ),142 "slow_query": (143 "SELECT *\n"144 "FROM users\n"145 "WHERE is_active = 1"146 ),147 },148 "subquery_to_join": {149 "name": "Correlated Subquery to JOIN",150 "difficulty": "medium",151 "max_steps": 6,152 "description": (153 "The slow query uses IN (SELECT ...) to filter orders, causing repeated inner-query execution. "154 "Rewrite it as an efficient JOIN. "155 "Return: order_id, user_id, total_amount "156 "for DELIVERED orders from ACTIVE users in the USA."157 ),158 "slow_query": (159 "SELECT order_id, user_id, total_amount\n"160 "FROM orders\n"161 "WHERE user_id IN (\n"162 " SELECT user_id\n"163 " FROM users\n"164 " WHERE country = 'USA'\n"165 " AND is_active = 1\n"166 ") AND status = 'delivered'"167 ),168 },169 "aggregation_optimization": {170 "name": "Aggregation Optimization",171 "difficulty": "hard",172 "max_steps": 8,173 "description": (174 "This query computes per-category revenue using correlated subqueries "175 "inside the SELECT list, which is O(n^2). "176 "Rewrite it using explicit JOINs and GROUP BY. "177 "Return: category_name (TEXT), total_revenue (REAL) "178 "for categories where total_revenue > 10000, "179 "ordered by total_revenue DESC."180 ),181 "slow_query": (182 "SELECT\n"183 " (SELECT name FROM categories\n"184 " WHERE category_id = p.category_id) AS category_name,\n"185 " (SELECT SUM(oi.quantity * oi.unit_price)\n"186 " FROM order_items oi\n"187 " WHERE oi.product_id IN (\n"188 " SELECT product_id FROM products\n"189 " WHERE category_id = p.category_id\n"190 " )) AS total_revenue\n"191 "FROM products p\n"192 "GROUP BY p.category_id\n"193 "HAVING total_revenue > 10000\n"194 "ORDER BY total_revenue DESC"195 ),196 },197 "cte_refactoring": {198 "name": "CTE Refactoring",199 "difficulty": "expert",200 "max_steps": 10,201 "description": (202 "This query identifies the top 5 customers by lifetime spend, "203 "using deeply nested subqueries that are hard to read and slow to execute. "204 "Rewrite it using Common Table Expressions (WITH clauses) for clarity and performance. "205 "Return: username, email, total_spend (REAL), order_count (INTEGER) "206 "for the top 5 active customers by total_spend, ordered by total_spend DESC."207 ),208 "slow_query": (209 "SELECT username, email, total_spend, order_count\n"210 "FROM (\n"211 " SELECT\n"212 " (SELECT username FROM users WHERE user_id = o.user_id) AS username,\n"213 " (SELECT email FROM users WHERE user_id = o.user_id) AS email,\n"214 " SUM(o.total_amount) AS total_spend,\n"215 " COUNT(*) AS order_count\n"216 " FROM orders o\n"217 " WHERE o.user_id IN (\n"218 " SELECT user_id FROM users WHERE is_active = 1\n"219 " )\n"220 " GROUP BY o.user_id\n"221 ") ranked\n"222 "ORDER BY total_spend DESC\n"223 "LIMIT 5"224 ),225 },226}227 228# ── SQL utilities ─────────────────────────────────────────────────────────────229 230def get_query_metrics(query: str, conn: sqlite3.Connection) -> ExecutionMetrics:231 try:232 plan_rows = conn.execute(f"EXPLAIN QUERY PLAN {query}").fetchall()233 plan_texts = [str(r) for r in plan_rows]234 flat = " ".join(plan_texts).upper()235 return ExecutionMetrics(236 uses_index=(237 "USING INDEX" in flat or238 ("SEARCH" in flat and flat.count("SCAN TABLE") == 0)239 ),240 full_scan_count=flat.count("SCAN TABLE"),241 query_plan=plan_texts,242 )243 except Exception:244 return ExecutionMetrics()245 246 247def _norm(q: str) -> str:248 return re.sub(r"\s+", " ", q.strip().upper())249 250def _has_select_star(q): return bool(re.search(r"SELECT\s+\*", _norm(q)))251def _uses_join(q): return bool(re.search(r"\bJOIN\b", _norm(q)))252def _uses_in_sub(q): return bool(re.search(r"\bIN\s*\(\s*SELECT\b", _norm(q)))253def _uses_sel_sub(q): return bool(re.search(r"SELECT\s+\(\s*SELECT\b", _norm(q)))254def _uses_cte(q): return bool(re.search(r"^\s*WITH\b", q.strip(), re.IGNORECASE))255def _uses_groupby(q): return bool(re.search(r"\bGROUP\s+BY\b", _norm(q)))256 257def _run(query: str, conn: sqlite3.Connection):258 cur = conn.execute(query)259 cols = [d[0].lower() for d in cur.description]260 rows = cur.fetchall()261 return cols, rows262 263# ── Graders ───────────────────────────────────────────────────────────────────264 265GraderResult = Tuple[float, Dict[str, float], str]266 267 268def grade_task1(query: str, conn: sqlite3.Connection) -> GraderResult:269 bd = dict(validity=0.0, correctness=0.0, performance=0.0, style=0.0)270 msgs = []271 272 try:273 cols, rows = _run(query, conn)274 bd["validity"] = 0.10275 msgs.append("✓ Valid SQL")276 except Exception as e:277 msgs.append(f"✗ SQL error: {e}")278 return 0.001, bd, " | ".join(msgs)279 280 required = ["user_id", "username", "email"]281 missing = [c for c in required if c not in cols]282 if missing:283 bd["correctness"] = 0.05284 msgs.append(f"✗ Missing columns: {missing}")285 else:286 ref = {r[0] for r in conn.execute("SELECT user_id FROM users WHERE is_active=1").fetchall()}287 got = {row[cols.index("user_id")] for row in rows}288 if got == ref:289 bd["correctness"] = 0.40290 msgs.append(f"✓ Correct ({len(ref)} users)")291 elif len(got & ref) / max(len(ref), 1) >= 0.95:292 bd["correctness"] = 0.20293 msgs.append("~ Partial correctness")294 else:295 msgs.append("✗ Wrong result set")296 297 if not _has_select_star(query):298 bd["style"] = 0.30299 msgs.append("✓ No SELECT *")300 else:301 msgs.append("✗ Still uses SELECT *")302 303 m = get_query_metrics(query, conn)304 if m.uses_index:305 bd["performance"] = 0.20306 msgs.append("✓ Index scan")307 elif m.full_scan_count <= 1:308 bd["performance"] = 0.10309 msgs.append("~ Single scan")310 else:311 msgs.append("✗ Multiple full scans")312 313 return min(0.999, max(0.001, sum(bd.values()))), bd, " | ".join(msgs)314 315 316def grade_task2(query: str, conn: sqlite3.Connection) -> GraderResult:317 bd = dict(validity=0.0, correctness=0.0, performance=0.0, style=0.0)318 msgs = []319 320 try:321 cols, rows = _run(query, conn)322 bd["validity"] = 0.10323 msgs.append("✓ Valid SQL")324 except Exception as e:325 msgs.append(f"✗ SQL error: {e}")326 return 0.001, bd, " | ".join(msgs)327 328 ref = {r[0] for r in conn.execute(329 "SELECT order_id FROM orders WHERE status='delivered' "330 "AND user_id IN (SELECT user_id FROM users WHERE country='USA' AND is_active=1)"331 ).fetchall()}332 333 required = ["order_id", "user_id", "total_amount"]334 missing = [c for c in required if c not in cols]335 if missing:336 bd["correctness"] = 0.05337 msgs.append(f"✗ Missing columns: {missing}")338 else:339 got = {row[cols.index("order_id")] for row in rows}340 overlap = len(got & ref) / max(len(ref), 1)341 if got == ref:342 bd["correctness"] = 0.40343 msgs.append(f"✓ Correct ({len(ref)} orders)")344 elif overlap >= 0.90:345 bd["correctness"] = 0.20346 msgs.append(f"~ {overlap:.0%} correct")347 else:348 msgs.append(f"✗ Wrong results ({overlap:.0%})")349 350 if not _uses_in_sub(query):351 bd["style"] = 0.20352 msgs.append("✓ No IN subquery")353 else:354 msgs.append("✗ Still uses IN subquery")355 356 uses_j = _uses_join(query)357 m = get_query_metrics(query, conn)358 if uses_j and m.uses_index:359 bd["performance"] = 0.30360 msgs.append("✓ JOIN + index")361 elif uses_j:362 bd["performance"] = 0.15363 msgs.append("~ JOIN but no index")364 elif m.uses_index:365 bd["performance"] = 0.10366 msgs.append("~ Index but no JOIN")367 else:368 msgs.append("✗ No JOIN, no index")369 370 return min(0.999, max(0.001, sum(bd.values()))), bd, " | ".join(msgs)371 372 373def grade_task3(query: str, conn: sqlite3.Connection) -> GraderResult:374 bd = dict(validity=0.0, correctness=0.0, performance=0.0, style=0.0)375 msgs = []376 377 try:378 cols, rows = _run(query, conn)379 bd["validity"] = 0.10380 msgs.append("✓ Valid SQL")381 except Exception as e:382 msgs.append(f"✗ SQL error: {e}")383 return 0.001, bd, " | ".join(msgs)384 385 ref_rows = conn.execute("""386 SELECT c.name, SUM(oi.quantity * oi.unit_price) AS total_revenue387 FROM categories c388 JOIN products p ON p.category_id = c.category_id389 JOIN order_items oi ON oi.product_id = p.product_id390 GROUP BY c.category_id, c.name391 HAVING SUM(oi.quantity * oi.unit_price) > 10000392 ORDER BY total_revenue DESC393 """).fetchall()394 ref_set = {(r[0], round(float(r[1]), 0)) for r in ref_rows}395 396 required = ["category_name", "total_revenue"]397 missing = [c for c in required if c not in cols]398 if missing:399 bd["correctness"] = 0.05400 msgs.append(f"✗ Missing columns: {missing}")401 else:402 got_set = {(row[cols.index("category_name")], round(float(row[cols.index("total_revenue")]), 0)) for row in rows}403 overlap = len(got_set & ref_set) / max(len(ref_set), 1)404 if got_set == ref_set:405 bd["correctness"] = 0.35406 msgs.append(f"✓ Exact match ({len(ref_set)} categories)")407 elif overlap >= 0.80:408 bd["correctness"] = 0.18409 msgs.append(f"~ {overlap:.0%} match")410 else:411 msgs.append(f"✗ Wrong results ({overlap:.0%})")412 413 if not _uses_sel_sub(query) and not _uses_in_sub(query):414 bd["style"] = 0.25415 msgs.append("✓ No correlated subqueries")416 elif not _uses_sel_sub(query):417 bd["style"] = 0.12418 msgs.append("~ Partial: SELECT subquery removed")419 else:420 msgs.append("✗ Correlated subqueries in SELECT")421 422 perf = 0.0423 if _uses_join(query): perf += 0.12; msgs.append("✓ JOIN")424 if _uses_groupby(query): perf += 0.10; msgs.append("✓ GROUP BY")425 m = get_query_metrics(query, conn)426 if m.uses_index: perf += 0.08; msgs.append("✓ Index")427 bd["performance"] = perf428 429 return min(0.999, max(0.001, sum(bd.values()))), bd, " | ".join(msgs)430 431 432def grade_task4(query: str, conn: sqlite3.Connection) -> GraderResult:433 """434 CTE Refactoring grader.435 436 Weights:437 validity 0.10438 correctness 0.35 — exact top-5 username+spend match439 style 0.30 — uses WITH (CTE), no nested subqueries440 performance 0.25 — JOIN + index441 """442 bd = dict(validity=0.0, correctness=0.0, performance=0.0, style=0.0)443 msgs = []444 445 try:446 cols, rows = _run(query, conn)447 bd["validity"] = 0.10448 msgs.append("✓ Valid SQL")449 except Exception as e:450 msgs.append(f"✗ SQL error: {e}")451 return 0.001, bd, " | ".join(msgs)452 453 ref_rows = conn.execute("""454 WITH active_users AS (455 SELECT user_id, username, email FROM users WHERE is_active = 1456 ),457 user_spend AS (458 SELECT o.user_id,459 SUM(o.total_amount) AS total_spend,460 COUNT(*) AS order_count461 FROM orders o462 JOIN active_users au ON au.user_id = o.user_id463 GROUP BY o.user_id464 )465 SELECT au.username, au.email,466 us.total_spend, us.order_count467 FROM user_spend us468 JOIN active_users au ON au.user_id = us.user_id469 ORDER BY us.total_spend DESC470 LIMIT 5471 """).fetchall()472 ref_set = {(r[0], round(float(r[2]), 1)) for r in ref_rows}473 474 required = ["username", "email", "total_spend", "order_count"]475 missing = [c for c in required if c not in cols]476 if missing:477 bd["correctness"] = 0.05478 msgs.append(f"✗ Missing columns: {missing}")479 elif len(rows) != 5:480 bd["correctness"] = 0.10481 msgs.append(f"✗ Expected 5 rows, got {len(rows)}")482 else:483 got_set = {(row[cols.index("username")], round(float(row[cols.index("total_spend")]), 1)) for row in rows}484 overlap = len(got_set & ref_set) / max(len(ref_set), 1)485 if got_set == ref_set:486 bd["correctness"] = 0.35487 msgs.append("✓ Exact top-5 match")488 elif overlap >= 0.80:489 bd["correctness"] = 0.18490 msgs.append(f"~ {overlap:.0%} of top-5 correct")491 else:492 msgs.append(f"✗ Wrong top-5 ({overlap:.0%})")493 494 uses_cte = _uses_cte(query)495 has_nest_sub = _uses_sel_sub(query) or (query.upper().count("SELECT") > 2 and _uses_in_sub(query))496 497 if uses_cte and not has_nest_sub:498 bd["style"] = 0.30499 msgs.append("✓ CTE used, no nested subqueries")500 elif uses_cte:501 bd["style"] = 0.15502 msgs.append("~ CTE used but nested subqueries remain")503 else:504 msgs.append("✗ No CTE (WITH clause) found")505 506 perf = 0.0507 if _uses_join(query): perf += 0.12; msgs.append("✓ JOIN")508 if _uses_groupby(query): perf += 0.08; msgs.append("✓ GROUP BY")509 m = get_query_metrics(query, conn)510 if m.uses_index: perf += 0.05; msgs.append("✓ Index")511 bd["performance"] = perf512 513 return min(0.999, max(0.001, sum(bd.values()))), bd, " | ".join(msgs)514 515 516# ── Grader registry ───────────────────────────────────────────────────────────517 518TASK_GRADERS = {519 "select_star_removal": grade_task1,520 "subquery_to_join": grade_task2,521 "aggregation_optimization": grade_task3,522 "cte_refactoring": grade_task4,523}524 525def _safe(fn):526 def w(q,conn):527 s,b,f=fn(q,conn)528 s=max(0.001,min(0.999,float(s)))529 b={k:max(0.001,min(0.999,float(v))) for k,v in b.items()}530 return s,b,f531 return w532TASK_GRADERS={k:_safe(v) for k,v in TASK_GRADERS.items()}533