File size: 5,599 Bytes
f72291c a575fa1 f72291c | 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 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | 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)) |