File size: 5,846 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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""
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)