mxpic_EDA/backend/database.py

# -----------------------------------------------------------------------------
# Description: SQLite database initialization and persistence helpers for users, projects, cells, and project metadata.
# Inside functions: connect_db, init_db, get_user, get_user_profile, get_user_auth_by_id, update_user_occupation, update_user_password, add_user_log, list_user_logs
# Developer : Qin Yue @ 2026
# Organization : OptiHK Limited
# -----------------------------------------------------------------------------
# backend/database.py
import sqlite3
import os
from werkzeug.security import generate_password_hash
from datetime import datetime

# Save the database in the backend folder
DB_FILE = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "database", "mxpic_data.db"))

def connect_db():
    conn = sqlite3.connect(DB_FILE, timeout=20)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=20000")
    return conn

def init_db():
    os.makedirs(os.path.dirname(DB_FILE), exist_ok=True)
    conn = connect_db()
    cursor = conn.cursor()

    # Create Users Table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL
        )
    ''')

    cursor.execute('''
        CREATE TABLE IF NOT EXISTS user_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            username TEXT NOT NULL,
            action TEXT NOT NULL,
            project TEXT,
            cell TEXT,
            detail TEXT,
            ip_address TEXT,
            created_at TEXT NOT NULL,
            FOREIGN KEY(user_id) REFERENCES users(id)
        )
    ''')

    cursor.execute("PRAGMA table_info(users)")
    existing_columns = {row[1] for row in cursor.fetchall()}
    migrations = {
        "created_at": "ALTER TABLE users ADD COLUMN created_at TEXT",
        "credits": "ALTER TABLE users ADD COLUMN credits INTEGER NOT NULL DEFAULT 0",
        "occupation": "ALTER TABLE users ADD COLUMN occupation TEXT NOT NULL DEFAULT 'intern'",
        "user_group": "ALTER TABLE users ADD COLUMN user_group TEXT NOT NULL DEFAULT 'user'",
    }
    for column, statement in migrations.items():
        if column not in existing_columns:
            cursor.execute(statement)

    now = datetime.utcnow().strftime("%Y-%m-%d")
    cursor.execute("UPDATE users SET created_at = ? WHERE created_at IS NULL OR created_at = ''", (now,))
    cursor.execute("UPDATE users SET user_group = 'manager' WHERE username = 'admin'")
    cursor.execute("UPDATE users SET user_group = 'developers' WHERE username = 'engineer'")
    cursor.execute("UPDATE users SET user_group = 'user' WHERE user_group IS NULL OR user_group = ''")
    
    # Insert default users for local multi-account development.
    cursor.execute("SELECT * FROM users WHERE username = 'admin'")
    if not cursor.fetchone():
        test_hash = generate_password_hash("123456")
        cursor.execute(
            "INSERT INTO users (username, password_hash, created_at, credits, occupation, user_group) VALUES (?, ?, ?, ?, ?, ?)",
            ("admin", test_hash, now, 0, "principle engineer", "manager")
        )
        print("Test user created. Username: admin | Password: 123456")

    cursor.execute("SELECT * FROM users WHERE username = 'engineer'")
    if not cursor.fetchone():
        test_hash = generate_password_hash("123456")
        cursor.execute(
            "INSERT INTO users (username, password_hash, created_at, credits, occupation, user_group) VALUES (?, ?, ?, ?, ?, ?)",
            ("engineer", test_hash, now, 0, "junior engineer", "developers")
        )
        print("Second test user created. Username: engineer | Password: 123456")

    conn.commit()
    conn.close()

def get_user(username):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute("SELECT id, username, password_hash, user_group FROM users WHERE username = ?", (username,))
    user = cursor.fetchone()
    conn.close()
    return user

def get_user_profile(user_id):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute(
        "SELECT id, username, created_at, credits, occupation, user_group FROM users WHERE id = ?",
        (user_id,)
    )
    user = cursor.fetchone()
    conn.close()
    return user

def get_user_auth_by_id(user_id):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute("SELECT id, username, password_hash FROM users WHERE id = ?", (user_id,))
    user = cursor.fetchone()
    conn.close()
    return user

def update_user_occupation(user_id, occupation):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute("UPDATE users SET occupation = ? WHERE id = ?", (occupation, user_id))
    conn.commit()
    conn.close()

def update_user_password(user_id, password):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute("UPDATE users SET password_hash = ? WHERE id = ?", (generate_password_hash(password), user_id))
    conn.commit()
    conn.close()

def add_user_log(user_id, username, action, project=None, cell=None, detail=None, ip_address=None):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute(
        '''
        INSERT INTO user_logs (user_id, username, action, project, cell, detail, ip_address, created_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''',
        (
            user_id,
            username,
            action,
            project,
            cell,
            detail,
            ip_address,
            datetime.utcnow().isoformat(timespec="seconds") + "Z"
        )
    )
    conn.commit()
    conn.close()

def list_user_logs(user_id, limit=200):
    conn = connect_db()
    cursor = conn.cursor()
    cursor.execute(
        '''
        SELECT id, action, project, cell, detail, ip_address, created_at
        FROM user_logs
        WHERE user_id = ?
        ORDER BY id DESC
        LIMIT ?
        ''',
        (user_id, limit)
    )
    rows = cursor.fetchall()
    conn.close()
    return rows

if __name__ == "__main__":
    init_db()
    print("Database initialized successfully.")