db.py

import sqlite3
from pathlib import Path

DB_PATH = Path("employees.db")

def get_connection():
    return sqlite3.connect(DB_PATH)

def init_db():
    """Create database and tables if they don't exist."""
    conn = get_connection()
    cursor = conn.cursor()

    # Employees table
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS employees (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            department TEXT NOT NULL,
            title TEXT
        )
    """)

    # Salaries table (FK to employees.id)
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS salaries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            employee_id INTEGER NOT NULL,
            salary REAL NOT NULL,
            FOREIGN KEY (employee_id) REFERENCES employees(id)
        )
    """)

    conn.commit()
    conn.close()


def fill_db():
    conn = get_connection()
    cursor = conn.cursor()

    # Seed employees
    cursor.execute("SELECT COUNT(*) FROM employees")
    emp_count = cursor.fetchone()[0]

    if emp_count == 0:
        cursor.executemany("""
            INSERT INTO employees (name, department, title)
            VALUES (?, ?, ?)
        """, [
            ("Alice Johnson", "Sales", "Account Manager"),
            ("Bob Smith", "Sales", "Sales Rep"),
            ("Carol White", "Engineering", "Backend Developer"),
            ("David Brown", "Engineering", "Frontend Developer"),
            ("Eve Davis", "HR", "HR Manager")
        ])

    # Seed salaries
    cursor.execute("SELECT COUNT(*) FROM salaries")
    sal_count = cursor.fetchone()[0]

    if sal_count == 0:
        cursor.executemany("""
            INSERT INTO salaries (employee_id, salary)
            VALUES (?, ?)
        """, [
            (1, 75000),
            (2, 65000),
            (3, 95000),
            (4, 90000),
            (5, 70000)
        ])

    conn.commit()
    conn.close()


def query_db(llm_query):
    conn = get_connection()
    cursor = conn.cursor()

    cursor.execute(llm_query)
    rows = cursor.fetchall()
    columns = [description[0] for description in cursor.description]
    results = [dict(zip(columns, row)) for row in rows]

    conn.close()

    return results