database.py

import sqlite3
import hashlib
import secrets
from datetime import datetime, timedelta
from typing import Optional, List, Dict
from contextlib import contextmanager
from config import settings


@contextmanager
def get_db():
    conn = sqlite3.connect(settings.database_url.replace("sqlite:///", ""))
    conn.row_factory = sqlite3.Row
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def init_db():
    with get_db() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                email TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                api_key TEXT UNIQUE NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                is_active BOOLEAN DEFAULT 1
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS download_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                session_id TEXT UNIQUE NOT NULL,
                status TEXT DEFAULT 'pending',
                ollama_version TEXT,
                model_name TEXT,
                progress INTEGER DEFAULT 0,
                current_step TEXT,
                total_size_mb INTEGER DEFAULT 0,
                downloaded_mb INTEGER DEFAULT 0,
                speed_kbps INTEGER DEFAULT 0,
                eta_seconds INTEGER DEFAULT 0,
                error_message TEXT,
                started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                completed_at TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS api_keys (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                key_name TEXT NOT NULL,
                api_key TEXT UNIQUE NOT NULL,
                is_active BOOLEAN DEFAULT 1,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_used_at TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        """)


def hash_password(password: str) -> str:
    return hashlib.sha256(password.encode()).hexdigest()


def verify_password(password: str, password_hash: str) -> bool:
    return hash_password(password) == password_hash


def generate_api_key() -> str:
    return f"ollama_{secrets.token_urlsafe(32)}"


def create_user(username: str, email: str, password: str) -> Optional[int]:
    with get_db() as conn:
        try:
            cursor = conn.execute(
                "INSERT INTO users (username, email, password_hash, api_key) VALUES (?, ?, ?, ?)",
                (username, email, hash_password(password), generate_api_key())
            )
            return cursor.lastrowid
        except sqlite3.IntegrityError:
            return None


def get_user_by_username(username: str) -> Optional[Dict]:
    with get_db() as conn:
        row = conn.execute("SELECT * FROM users WHERE username = ?", (username,)).fetchone()
        return dict(row) if row else None


def get_user_by_api_key(api_key: str) -> Optional[Dict]:
    with get_db() as conn:
        row = conn.execute("SELECT * FROM users WHERE api_key = ? AND is_active = 1", (api_key,)).fetchone()
        if row:
            return dict(row)
        row = conn.execute(
            "SELECT u.* FROM users u JOIN api_keys a ON u.id = a.user_id WHERE a.api_key = ? AND a.is_active = 1",
            (api_key,)
        ).fetchone()
        return dict(row) if row else None


def create_download_session(user_id: int, session_id: str, model_name: str) -> int:
    with get_db() as conn:
        cursor = conn.execute(
            "INSERT INTO download_sessions (user_id, session_id, model_name, status) VALUES (?, ?, ?, 'pending')",
            (user_id, session_id, model_name)
        )
        return cursor.lastrowid


def update_download_progress(session_id: str, **kwargs):
    with get_db() as conn:
        fields = ", ".join([f"{k} = ?" for k in kwargs.keys()])
        values = list(kwargs.values()) + [session_id]
        conn.execute(f"UPDATE download_sessions SET {fields} WHERE session_id = ?", values)


def get_download_session(session_id: str) -> Optional[Dict]:
    with get_db() as conn:
        row = conn.execute("SELECT * FROM download_sessions WHERE session_id = ?", (session_id,)).fetchone()
        return dict(row) if row else None


def get_user_sessions(user_id: int) -> List[Dict]:
    with get_db() as conn:
        rows = conn.execute(
            "SELECT * FROM download_sessions WHERE user_id = ? ORDER BY started_at DESC",
            (user_id,)
        ).fetchall()
        return [dict(row) for row in rows]


def create_api_key(user_id: int, key_name: str) -> str:
    api_key = generate_api_key()
    with get_db() as conn:
        conn.execute(
            "INSERT INTO api_keys (user_id, key_name, api_key) VALUES (?, ?, ?)",
            (user_id, key_name, api_key)
        )
    return api_key


def get_user_api_keys(user_id: int) -> List[Dict]:
    with get_db() as conn:
        rows = conn.execute("SELECT * FROM api_keys WHERE user_id = ?", (user_id,)).fetchall()
        return [dict(row) for row in rows]


def revoke_api_key(key_id: int, user_id: int):
    with get_db() as conn:
        conn.execute("UPDATE api_keys SET is_active = 0 WHERE id = ? AND user_id = ?", (key_id, user_id))