server/database.py

"""
Database setup for the SQL Query Environment.

Creates an in-memory SQLite database with three tables:
  - departments: id, name, budget, location
  - employees: id, name, department_id, salary, hire_date, is_active
  - projects: id, name, department_id, lead_employee_id, budget, status, start_date

All data is deterministic so grading is reproducible.
"""

import sqlite3
from typing import Optional

SCHEMA_SQL = """
CREATE TABLE departments (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    budget REAL NOT NULL,
    location TEXT NOT NULL
);

CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    department_id INTEGER NOT NULL,
    salary REAL NOT NULL,
    hire_date TEXT NOT NULL,
    is_active INTEGER NOT NULL DEFAULT 1,
    FOREIGN KEY (department_id) REFERENCES departments(id)
);

CREATE TABLE projects (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    department_id INTEGER NOT NULL,
    lead_employee_id INTEGER NOT NULL,
    budget REAL NOT NULL,
    status TEXT NOT NULL CHECK(status IN ('active', 'completed', 'cancelled')),
    start_date TEXT NOT NULL,
    FOREIGN KEY (department_id) REFERENCES departments(id),
    FOREIGN KEY (lead_employee_id) REFERENCES employees(id)
);
"""

SEED_SQL = """
INSERT INTO departments VALUES (1, 'Engineering', 500000.00, 'Bangalore');
INSERT INTO departments VALUES (2, 'Marketing', 200000.00, 'Mumbai');
INSERT INTO departments VALUES (3, 'Sales', 300000.00, 'Delhi');
INSERT INTO departments VALUES (4, 'HR', 150000.00, 'Bangalore');
INSERT INTO departments VALUES (5, 'Finance', 250000.00, 'Mumbai');

INSERT INTO employees VALUES (1,  'Arjun Sharma',    1, 85000.00, '2020-03-15', 1);
INSERT INTO employees VALUES (2,  'Priya Patel',     1, 92000.00, '2019-07-01', 1);
INSERT INTO employees VALUES (3,  'Rahul Verma',     1, 78000.00, '2021-01-10', 1);
INSERT INTO employees VALUES (4,  'Sneha Gupta',     2, 65000.00, '2020-06-20', 1);
INSERT INTO employees VALUES (5,  'Vikram Singh',    2, 70000.00, '2018-11-05', 1);
INSERT INTO employees VALUES (6,  'Anita Desai',     3, 72000.00, '2019-09-12', 1);
INSERT INTO employees VALUES (7,  'Karan Mehta',     3, 68000.00, '2021-04-01', 1);
INSERT INTO employees VALUES (8,  'Deepa Nair',      3, 75000.00, '2020-02-28', 0);
INSERT INTO employees VALUES (9,  'Suresh Kumar',    4, 60000.00, '2022-01-15', 1);
INSERT INTO employees VALUES (10, 'Meera Joshi',     4, 58000.00, '2021-08-20', 1);
INSERT INTO employees VALUES (11, 'Amit Rao',        5, 88000.00, '2019-05-10', 1);
INSERT INTO employees VALUES (12, 'Lakshmi Iyer',    5, 82000.00, '2020-10-01', 1);
INSERT INTO employees VALUES (13, 'Ravi Krishnan',   1, 95000.00, '2018-03-20', 1);
INSERT INTO employees VALUES (14, 'Pooja Reddy',     2, 62000.00, '2022-06-15', 1);
INSERT INTO employees VALUES (15, 'Nikhil Agarwal',  3, 71000.00, '2020-12-01', 1);

INSERT INTO projects VALUES (1, 'Cloud Migration',     1, 2,  120000.00, 'active',    '2024-01-15');
INSERT INTO projects VALUES (2, 'Mobile App v2',       1, 1,  80000.00,  'active',    '2024-03-01');
INSERT INTO projects VALUES (3, 'Brand Refresh',       2, 5,  45000.00,  'completed', '2023-06-01');
INSERT INTO projects VALUES (4, 'Q4 Campaign',         2, 4,  60000.00,  'active',    '2024-09-01');
INSERT INTO projects VALUES (5, 'CRM Integration',     3, 6,  90000.00,  'active',    '2024-02-15');
INSERT INTO projects VALUES (6, 'Sales Dashboard',     3, 15, 35000.00,  'completed', '2023-11-01');
INSERT INTO projects VALUES (7, 'Payroll Automation',   4, 9,  50000.00,  'active',    '2024-04-01');
INSERT INTO projects VALUES (8, 'Annual Audit Tool',    5, 11, 70000.00,  'cancelled', '2024-01-10');
INSERT INTO projects VALUES (9, 'Data Pipeline',        1, 13, 150000.00, 'active',    '2024-06-01');
INSERT INTO projects VALUES (10, 'Employee Portal',     4, 10, 40000.00,  'completed', '2023-09-15');
"""

SCHEMA_DESCRIPTION = """Tables in the database:

CREATE TABLE departments (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    budget REAL NOT NULL,
    location TEXT NOT NULL
);

CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    department_id INTEGER NOT NULL,
    salary REAL NOT NULL,
    hire_date TEXT NOT NULL,       -- format: YYYY-MM-DD
    is_active INTEGER NOT NULL,    -- 1 = active, 0 = inactive
    FOREIGN KEY (department_id) REFERENCES departments(id)
);

CREATE TABLE projects (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    department_id INTEGER NOT NULL,
    lead_employee_id INTEGER NOT NULL,
    budget REAL NOT NULL,
    status TEXT NOT NULL,           -- 'active', 'completed', or 'cancelled'
    start_date TEXT NOT NULL,       -- format: YYYY-MM-DD
    FOREIGN KEY (department_id) REFERENCES departments(id),
    FOREIGN KEY (lead_employee_id) REFERENCES employees(id)
);""".strip()


def create_database() -> sqlite3.Connection:
    """Create a fresh in-memory SQLite database with schema and seed data."""
    conn = sqlite3.connect(":memory:")
    conn.execute("PRAGMA foreign_keys = ON;")
    conn.executescript(SCHEMA_SQL)
    conn.executescript(SEED_SQL)
    conn.commit()
    return conn


def execute_query(
    conn: sqlite3.Connection, sql: str
) -> tuple[list[tuple] | None, list[str] | None, str | None]:
    """
    Execute a SQL query safely.

    Returns:
        (rows, column_names, error_message)
        - On success: (rows_list, columns_list, None)
        - On error: (None, None, error_string)
    """
    try:
        cursor = conn.execute(sql)
        if cursor.description is not None:
            columns = [desc[0] for desc in cursor.description]
            rows = cursor.fetchall()
            return rows, columns, None
        else:
            conn.commit()
            return None, None, None
    except Exception as e:
        return None, None, str(e)