File size: 4,670 Bytes
ede2fa4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
"""
Task definitions for the SQL Query Environment.

Each task has:
  - task_id: unique identifier
  - description: natural language question
  - difficulty: easy / medium / hard
  - expected_sql: a reference SQL that produces the correct answer
  - expected_columns: the column names in the expected result
  - expected_rows: the exact rows expected (deterministic)
"""

from dataclasses import dataclass, field


@dataclass
class TaskDefinition:
    task_id: str
    description: str
    difficulty: str
    expected_sql: str
    expected_columns: list[str]
    expected_rows: list[tuple]


# ──────────────────────────────────────────────
# TASK 1 β€” EASY: Single table, simple WHERE + ORDER BY
# ──────────────────────────────────────────────
TASK_1 = TaskDefinition(
    task_id="task_1",
    description=(
        "List the names and salaries of all active employees who earn more "
        "than 70000, ordered by salary descending."
    ),
    difficulty="easy",
    expected_sql="""
        SELECT name, salary
        FROM employees
        WHERE is_active = 1 AND salary > 70000
        ORDER BY salary DESC;
    """,
    expected_columns=["name", "salary"],
    expected_rows=[
        ("Ravi Krishnan", 95000.0),
        ("Priya Patel", 92000.0),
        ("Amit Rao", 88000.0),
        ("Arjun Sharma", 85000.0),
        ("Lakshmi Iyer", 82000.0),
        ("Rahul Verma", 78000.0),
        ("Anita Desai", 72000.0),
        ("Nikhil Agarwal", 71000.0),
    ],
)

# ──────────────────────────────────────────────
# TASK 2 β€” MEDIUM: JOIN + GROUP BY + HAVING
# ──────────────────────────────────────────────
TASK_2 = TaskDefinition(
    task_id="task_2",
    description=(
        "For each department, show the department name, the number of active "
        "employees, and the average salary of active employees. Only include "
        "departments that have more than 2 active employees. "
        "Order by average salary descending."
    ),
    difficulty="medium",
    expected_sql="""
        SELECT d.name, COUNT(e.id) AS num_employees, ROUND(AVG(e.salary), 2) AS avg_salary
        FROM departments d
        JOIN employees e ON d.id = e.department_id
        WHERE e.is_active = 1
        GROUP BY d.id, d.name
        HAVING COUNT(e.id) > 2
        ORDER BY avg_salary DESC;
    """,
    expected_columns=["name", "num_employees", "avg_salary"],
    expected_rows=[
        ("Engineering", 4, 87500.0),
        ("Sales", 3, 70333.33),
        ("Marketing", 3, 65666.67),
    ],
)

# ──────────────────────────────────────────────
# TASK 3 β€” HARD: Subquery + multiple JOINs + complex logic
# ──────────────────────────────────────────────
TASK_3 = TaskDefinition(
    task_id="task_3",
    description=(
        "Find all active employees who lead at least one active project and "
        "whose salary is above the overall average salary of all active employees. "
        "For each such employee, show their name, salary, department name, and "
        "the total budget of active projects they lead. "
        "Order by total project budget descending."
    ),
    difficulty="hard",
    expected_sql="""
        SELECT
            e.name,
            e.salary,
            d.name AS department_name,
            SUM(p.budget) AS total_project_budget
        FROM employees e
        JOIN departments d ON e.department_id = d.id
        JOIN projects p ON p.lead_employee_id = e.id
        WHERE e.is_active = 1
          AND p.status = 'active'
          AND e.salary > (
              SELECT AVG(salary) FROM employees WHERE is_active = 1
          )
        GROUP BY e.id, e.name, e.salary, d.name
        ORDER BY total_project_budget DESC;
    """,
    expected_columns=["name", "salary", "department_name", "total_project_budget"],
    expected_rows=[
        ("Ravi Krishnan", 95000.0, "Engineering", 150000.0),
        ("Priya Patel", 92000.0, "Engineering", 120000.0),
        ("Arjun Sharma", 85000.0, "Engineering", 80000.0),
    ],
)

# All tasks in order
ALL_TASKS = {
    "task_1": TASK_1,
    "task_2": TASK_2,
    "task_3": TASK_3,
}

TASK_LIST = [TASK_1, TASK_2, TASK_3]