server/tasks.py

"""
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]