| from __future__ import annotations |
|
|
| import json |
| import sqlite3 |
| import uuid |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| TERMINAL_TASK_STATUSES = {'success', 'failed', 'stopped', 'interrupted'} |
| ACTIVE_TASK_STATUSES = {'queued', 'running', 'waiting_captcha'} |
|
|
|
|
| def utc_now() -> str: |
| return datetime.now(timezone.utc).isoformat(timespec='seconds') |
|
|
|
|
| class Database: |
| def __init__(self, db_path: Path): |
| self.db_path = Path(db_path) |
| self.db_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| def _connect(self) -> sqlite3.Connection: |
| connection = sqlite3.connect(self.db_path, check_same_thread=False) |
| connection.row_factory = sqlite3.Row |
| connection.execute('PRAGMA foreign_keys = ON') |
| return connection |
|
|
| def initialize(self) -> None: |
| with self._connect() as connection: |
| connection.execute('PRAGMA journal_mode = WAL') |
| connection.executescript( |
| """ |
| CREATE TABLE IF NOT EXISTS users ( |
| id INTEGER PRIMARY KEY AUTOINCREMENT, |
| student_id TEXT NOT NULL UNIQUE, |
| display_name TEXT NOT NULL DEFAULT '', |
| encrypted_password TEXT NOT NULL, |
| created_at TEXT NOT NULL, |
| updated_at TEXT NOT NULL |
| ); |
| |
| CREATE TABLE IF NOT EXISTS admins ( |
| id INTEGER PRIMARY KEY AUTOINCREMENT, |
| username TEXT NOT NULL UNIQUE, |
| password_hash TEXT NOT NULL, |
| created_at TEXT NOT NULL, |
| updated_at TEXT NOT NULL |
| ); |
| |
| CREATE TABLE IF NOT EXISTS courses ( |
| id INTEGER PRIMARY KEY AUTOINCREMENT, |
| user_id INTEGER NOT NULL, |
| course_id TEXT NOT NULL, |
| course_index TEXT NOT NULL, |
| created_at TEXT NOT NULL, |
| UNIQUE(user_id, course_id, course_index), |
| FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE |
| ); |
| |
| CREATE TABLE IF NOT EXISTS tasks ( |
| id TEXT PRIMARY KEY, |
| user_id INTEGER NOT NULL, |
| status TEXT NOT NULL, |
| created_at TEXT NOT NULL, |
| started_at TEXT, |
| finished_at TEXT, |
| requested_by_role TEXT NOT NULL, |
| requested_by_identity TEXT NOT NULL, |
| stop_requested INTEGER NOT NULL DEFAULT 0, |
| last_error TEXT NOT NULL DEFAULT '', |
| total_count INTEGER NOT NULL DEFAULT 0, |
| completed_count INTEGER NOT NULL DEFAULT 0, |
| task_payload TEXT NOT NULL, |
| FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE |
| ); |
| |
| CREATE TABLE IF NOT EXISTS task_logs ( |
| id INTEGER PRIMARY KEY AUTOINCREMENT, |
| task_id TEXT NOT NULL, |
| level TEXT NOT NULL, |
| message TEXT NOT NULL, |
| created_at TEXT NOT NULL, |
| FOREIGN KEY(task_id) REFERENCES tasks(id) ON DELETE CASCADE |
| ); |
| |
| CREATE TABLE IF NOT EXISTS settings ( |
| key TEXT PRIMARY KEY, |
| value TEXT NOT NULL |
| ); |
| |
| CREATE INDEX IF NOT EXISTS idx_courses_user_id ON courses(user_id); |
| CREATE INDEX IF NOT EXISTS idx_tasks_user_id ON tasks(user_id); |
| CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status); |
| CREATE INDEX IF NOT EXISTS idx_task_logs_task_id ON task_logs(task_id, id); |
| """ |
| ) |
| connection.execute( |
| 'INSERT OR IGNORE INTO settings(key, value) VALUES (?, ?)', |
| ('max_parallel_tasks', '2'), |
| ) |
| connection.execute( |
| """ |
| UPDATE tasks |
| SET status = 'interrupted', |
| finished_at = ?, |
| last_error = CASE |
| WHEN COALESCE(last_error, '') = '' THEN '应用重启，上一轮任务被中断。' |
| ELSE last_error |
| END |
| WHERE status IN ('queued', 'running', 'waiting_captcha') |
| """, |
| (utc_now(),), |
| ) |
|
|
| def get_setting(self, key: str, default: str = '') -> str: |
| with self._connect() as connection: |
| row = connection.execute('SELECT value FROM settings WHERE key = ?', (key,)).fetchone() |
| return row['value'] if row else default |
|
|
| def set_setting(self, key: str, value: str) -> None: |
| with self._connect() as connection: |
| connection.execute( |
| """ |
| INSERT INTO settings(key, value) VALUES (?, ?) |
| ON CONFLICT(key) DO UPDATE SET value = excluded.value |
| """, |
| (key, value), |
| ) |
|
|
| def get_max_parallel_tasks(self) -> int: |
| raw = self.get_setting('max_parallel_tasks', '2') |
| try: |
| return max(1, min(8, int(raw))) |
| except ValueError: |
| return 2 |
|
|
| def create_user(self, student_id: str, encrypted_password: str, display_name: str = '') -> int: |
| now = utc_now() |
| with self._connect() as connection: |
| cursor = connection.execute( |
| """ |
| INSERT INTO users(student_id, display_name, encrypted_password, created_at, updated_at) |
| VALUES (?, ?, ?, ?, ?) |
| """, |
| (student_id, display_name, encrypted_password, now, now), |
| ) |
| return int(cursor.lastrowid) |
|
|
| def update_user( |
| self, |
| user_id: int, |
| *, |
| student_id: str | None = None, |
| encrypted_password: str | None = None, |
| display_name: str | None = None, |
| ) -> None: |
| fields: list[str] = [] |
| values: list[Any] = [] |
| if student_id is not None: |
| fields.append('student_id = ?') |
| values.append(student_id) |
| if encrypted_password is not None: |
| fields.append('encrypted_password = ?') |
| values.append(encrypted_password) |
| if display_name is not None: |
| fields.append('display_name = ?') |
| values.append(display_name) |
| if not fields: |
| return |
| fields.append('updated_at = ?') |
| values.append(utc_now()) |
| values.append(user_id) |
| with self._connect() as connection: |
| connection.execute(f"UPDATE users SET {', '.join(fields)} WHERE id = ?", values) |
|
|
| def delete_user(self, user_id: int) -> None: |
| with self._connect() as connection: |
| connection.execute('DELETE FROM users WHERE id = ?', (user_id,)) |
|
|
| def get_user_by_student_id(self, student_id: str) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute('SELECT * FROM users WHERE student_id = ?', (student_id,)).fetchone() |
| return dict(row) if row else None |
|
|
| def get_user_by_id(self, user_id: int) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute('SELECT * FROM users WHERE id = ?', (user_id,)).fetchone() |
| return dict(row) if row else None |
|
|
| def list_users(self) -> list[dict[str, Any]]: |
| with self._connect() as connection: |
| rows = connection.execute( |
| """ |
| SELECT |
| u.*, |
| COUNT(c.id) AS course_count, |
| ( |
| SELECT t.status |
| FROM tasks t |
| WHERE t.user_id = u.id |
| ORDER BY t.created_at DESC |
| LIMIT 1 |
| ) AS latest_task_status |
| FROM users u |
| LEFT JOIN courses c ON c.user_id = u.id |
| GROUP BY u.id |
| ORDER BY u.created_at ASC |
| """ |
| ).fetchall() |
| return [dict(row) for row in rows] |
|
|
| def list_courses_for_user(self, user_id: int) -> list[dict[str, Any]]: |
| with self._connect() as connection: |
| rows = connection.execute( |
| """ |
| SELECT id, course_id, course_index, created_at |
| FROM courses |
| WHERE user_id = ? |
| ORDER BY created_at ASC, id ASC |
| """, |
| (user_id,), |
| ).fetchall() |
| return [dict(row) for row in rows] |
|
|
| def add_course(self, user_id: int, course_id: str, course_index: str) -> None: |
| with self._connect() as connection: |
| connection.execute( |
| """ |
| INSERT OR IGNORE INTO courses(user_id, course_id, course_index, created_at) |
| VALUES (?, ?, ?, ?) |
| """, |
| (user_id, course_id, course_index, utc_now()), |
| ) |
|
|
| def delete_course(self, course_row_id: int, user_id: int | None = None) -> None: |
| with self._connect() as connection: |
| if user_id is None: |
| connection.execute('DELETE FROM courses WHERE id = ?', (course_row_id,)) |
| else: |
| connection.execute('DELETE FROM courses WHERE id = ? AND user_id = ?', (course_row_id, user_id)) |
|
|
| def list_admins(self) -> list[dict[str, Any]]: |
| with self._connect() as connection: |
| rows = connection.execute('SELECT id, username, created_at, updated_at FROM admins ORDER BY created_at ASC').fetchall() |
| return [dict(row) for row in rows] |
|
|
| def get_admin_by_username(self, username: str) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute('SELECT * FROM admins WHERE username = ?', (username,)).fetchone() |
| return dict(row) if row else None |
|
|
| def create_admin(self, username: str, password_hash: str) -> int: |
| now = utc_now() |
| with self._connect() as connection: |
| cursor = connection.execute( |
| """ |
| INSERT INTO admins(username, password_hash, created_at, updated_at) |
| VALUES (?, ?, ?, ?) |
| """, |
| (username, password_hash, now, now), |
| ) |
| return int(cursor.lastrowid) |
|
|
| def update_admin_password(self, admin_id: int, password_hash: str) -> None: |
| with self._connect() as connection: |
| connection.execute( |
| 'UPDATE admins SET password_hash = ?, updated_at = ? WHERE id = ?', |
| (password_hash, utc_now(), admin_id), |
| ) |
|
|
| def delete_admin(self, admin_id: int) -> None: |
| with self._connect() as connection: |
| connection.execute('DELETE FROM admins WHERE id = ?', (admin_id,)) |
|
|
| def find_active_task_for_user(self, user_id: int) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute( |
| """ |
| SELECT * FROM tasks |
| WHERE user_id = ? AND status IN ('queued', 'running', 'waiting_captcha') |
| ORDER BY created_at DESC |
| LIMIT 1 |
| """, |
| (user_id,), |
| ).fetchone() |
| return dict(row) if row else None |
|
|
| def create_task( |
| self, |
| *, |
| user_id: int, |
| requested_by_role: str, |
| requested_by_identity: str, |
| payload: dict[str, Any], |
| ) -> str: |
| task_id = str(uuid.uuid4()) |
| now = utc_now() |
| with self._connect() as connection: |
| connection.execute( |
| """ |
| INSERT INTO tasks( |
| id, |
| user_id, |
| status, |
| created_at, |
| requested_by_role, |
| requested_by_identity, |
| total_count, |
| completed_count, |
| task_payload |
| ) |
| VALUES (?, ?, 'queued', ?, ?, ?, ?, 0, ?) |
| """, |
| ( |
| task_id, |
| user_id, |
| now, |
| requested_by_role, |
| requested_by_identity, |
| len(payload.get('courses', [])), |
| json.dumps(payload, ensure_ascii=False), |
| ), |
| ) |
| return task_id |
|
|
| def claim_next_queued_task(self) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute("SELECT id FROM tasks WHERE status = 'queued' ORDER BY created_at ASC LIMIT 1").fetchone() |
| if not row: |
| return None |
| updated = connection.execute( |
| """ |
| UPDATE tasks |
| SET status = 'running', started_at = ?, last_error = '' |
| WHERE id = ? AND status = 'queued' |
| """, |
| (utc_now(), row['id']), |
| ).rowcount |
| if not updated: |
| return None |
| claimed = connection.execute('SELECT * FROM tasks WHERE id = ?', (row['id'],)).fetchone() |
| return dict(claimed) if claimed else None |
|
|
| def get_task(self, task_id: str) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute('SELECT * FROM tasks WHERE id = ?', (task_id,)).fetchone() |
| return dict(row) if row else None |
|
|
| def get_task_with_user(self, task_id: str) -> dict[str, Any] | None: |
| with self._connect() as connection: |
| row = connection.execute( |
| """ |
| SELECT |
| t.*, |
| u.student_id, |
| u.display_name |
| FROM tasks t |
| JOIN users u ON u.id = t.user_id |
| WHERE t.id = ? |
| """, |
| (task_id,), |
| ).fetchone() |
| return dict(row) if row else None |
|
|
| def list_recent_tasks_for_user(self, user_id: int, limit: int = 12) -> list[dict[str, Any]]: |
| with self._connect() as connection: |
| rows = connection.execute( |
| """ |
| SELECT id, user_id, status, created_at, started_at, finished_at, stop_requested, |
| last_error, total_count, completed_count |
| FROM tasks |
| WHERE user_id = ? |
| ORDER BY created_at DESC |
| LIMIT ? |
| """, |
| (user_id, limit), |
| ).fetchall() |
| return [dict(row) for row in rows] |
|
|
| def list_recent_tasks(self, limit: int = 20) -> list[dict[str, Any]]: |
| with self._connect() as connection: |
| rows = connection.execute( |
| """ |
| SELECT |
| t.id, |
| t.user_id, |
| t.status, |
| t.created_at, |
| t.started_at, |
| t.finished_at, |
| t.stop_requested, |
| t.last_error, |
| t.total_count, |
| t.completed_count, |
| u.student_id, |
| u.display_name |
| FROM tasks t |
| JOIN users u ON u.id = t.user_id |
| ORDER BY t.created_at DESC |
| LIMIT ? |
| """, |
| (limit,), |
| ).fetchall() |
| return [dict(row) for row in rows] |
|
|
| def set_task_status( |
| self, |
| task_id: str, |
| status: str, |
| *, |
| last_error: str | None = None, |
| completed_count: int | None = None, |
| ) -> None: |
| assignments = ['status = ?'] |
| values: list[Any] = [status] |
| if last_error is not None: |
| assignments.append('last_error = ?') |
| values.append(last_error) |
| if completed_count is not None: |
| assignments.append('completed_count = ?') |
| values.append(completed_count) |
| if status == 'running': |
| assignments.append('started_at = COALESCE(started_at, ?)') |
| values.append(utc_now()) |
| if status in TERMINAL_TASK_STATUSES: |
| assignments.append('finished_at = ?') |
| values.append(utc_now()) |
| values.append(task_id) |
| with self._connect() as connection: |
| connection.execute(f"UPDATE tasks SET {', '.join(assignments)} WHERE id = ?", values) |
|
|
| def update_task_progress(self, task_id: str, completed_count: int) -> None: |
| with self._connect() as connection: |
| connection.execute('UPDATE tasks SET completed_count = ? WHERE id = ?', (completed_count, task_id)) |
|
|
| def request_task_stop(self, task_id: str) -> None: |
| with self._connect() as connection: |
| connection.execute('UPDATE tasks SET stop_requested = 1 WHERE id = ?', (task_id,)) |
|
|
| def is_task_stop_requested(self, task_id: str) -> bool: |
| with self._connect() as connection: |
| row = connection.execute('SELECT stop_requested FROM tasks WHERE id = ?', (task_id,)).fetchone() |
| return bool(row and row['stop_requested']) |
|
|
| def append_task_log(self, task_id: str, level: str, message: str) -> int: |
| with self._connect() as connection: |
| cursor = connection.execute( |
| """ |
| INSERT INTO task_logs(task_id, level, message, created_at) |
| VALUES (?, ?, ?, ?) |
| """, |
| (task_id, level, message, utc_now()), |
| ) |
| return int(cursor.lastrowid) |
|
|
| def list_task_logs(self, task_id: str, after_id: int = 0, limit: int = 200) -> list[dict[str, Any]]: |
| with self._connect() as connection: |
| rows = connection.execute( |
| """ |
| SELECT id, level, message, created_at |
| FROM task_logs |
| WHERE task_id = ? AND id > ? |
| ORDER BY id ASC |
| LIMIT ? |
| """, |
| (task_id, after_id, limit), |
| ).fetchall() |
| return [dict(row) for row in rows] |
|
|
| def get_system_snapshot(self) -> dict[str, Any]: |
| with self._connect() as connection: |
| return { |
| 'users': connection.execute('SELECT COUNT(*) FROM users').fetchone()[0], |
| 'admins': connection.execute('SELECT COUNT(*) FROM admins').fetchone()[0], |
| 'queued': connection.execute("SELECT COUNT(*) FROM tasks WHERE status = 'queued'").fetchone()[0], |
| 'running': connection.execute( |
| "SELECT COUNT(*) FROM tasks WHERE status IN ('running', 'waiting_captcha')" |
| ).fetchone()[0], |
| 'max_parallel_tasks': self.get_max_parallel_tasks(), |
| }
|
|
|
