296 lines
8.7 KiB
Python
296 lines
8.7 KiB
Python
"""
|
|
Background Tasks & Async Job Manager.
|
|
100% pure Python stdlib + SQLite WAL mode.
|
|
Zero Redis, zero Celery, zero bloat (Ponytail architecture).
|
|
"""
|
|
|
|
import json
|
|
import sqlite3
|
|
import asyncio
|
|
from typing import Dict, Any, List, Optional
|
|
from datetime import datetime
|
|
|
|
# In-memory cancellation registry for instantaneous cooperative cancellation
|
|
CANCELLED_TASK_IDS = set()
|
|
|
|
|
|
def create_task(
|
|
conn: sqlite3.Connection,
|
|
user_id: int,
|
|
task_type: str,
|
|
params: Optional[Dict[str, Any]] = None,
|
|
total: int = 0,
|
|
message: str = "Задача поставлена в очередь",
|
|
) -> int:
|
|
"""
|
|
Creates a new background task in 'pending' status.
|
|
"""
|
|
params_json = json.dumps(params or {}, ensure_ascii=False)
|
|
cursor = conn.execute(
|
|
"""
|
|
INSERT INTO background_tasks (
|
|
user_id, task_type, status, progress_current, progress_total,
|
|
message, params_json, result_json, error_message,
|
|
created_at, updated_at
|
|
) VALUES (?, ?, 'pending', 0, ?, ?, ?, '{}', '', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
|
""",
|
|
(user_id, task_type, total, message, params_json),
|
|
)
|
|
return cursor.lastrowid
|
|
|
|
|
|
def update_task_progress(
|
|
conn: sqlite3.Connection,
|
|
task_id: int,
|
|
current: int,
|
|
total: Optional[int] = None,
|
|
message: Optional[str] = None,
|
|
):
|
|
"""
|
|
Updates the progress and status of a running task.
|
|
"""
|
|
if total is not None and message is not None:
|
|
conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'running',
|
|
progress_current = ?,
|
|
progress_total = ?,
|
|
message = ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND status != 'cancelled'
|
|
""",
|
|
(current, total, message, task_id),
|
|
)
|
|
elif total is not None:
|
|
conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'running',
|
|
progress_current = ?,
|
|
progress_total = ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND status != 'cancelled'
|
|
""",
|
|
(current, total, task_id),
|
|
)
|
|
elif message is not None:
|
|
conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'running',
|
|
progress_current = ?,
|
|
message = ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND status != 'cancelled'
|
|
""",
|
|
(current, message, task_id),
|
|
)
|
|
else:
|
|
conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'running',
|
|
progress_current = ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND status != 'cancelled'
|
|
""",
|
|
(current, task_id),
|
|
)
|
|
|
|
|
|
def complete_task(
|
|
conn: sqlite3.Connection,
|
|
task_id: int,
|
|
result: Optional[Dict[str, Any]] = None,
|
|
message: str = "Задача успешно завершена",
|
|
):
|
|
"""
|
|
Marks a background task as completed.
|
|
"""
|
|
result_json = json.dumps(result or {}, ensure_ascii=False)
|
|
conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'completed',
|
|
result_json = ?,
|
|
message = ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND status != 'cancelled'
|
|
""",
|
|
(result_json, message, task_id),
|
|
)
|
|
CANCELLED_TASK_IDS.discard(task_id)
|
|
|
|
|
|
def fail_task(
|
|
conn: sqlite3.Connection,
|
|
task_id: int,
|
|
error_message: str,
|
|
):
|
|
"""
|
|
Marks a background task as failed with an error message.
|
|
"""
|
|
conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'failed',
|
|
error_message = ?,
|
|
message = 'Ошибка выполнения задачи',
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND status != 'cancelled'
|
|
""",
|
|
(error_message, task_id),
|
|
)
|
|
CANCELLED_TASK_IDS.discard(task_id)
|
|
|
|
|
|
def cancel_task(conn: sqlite3.Connection, task_id: int, user_id: int) -> bool:
|
|
"""
|
|
Cooperatively cancels a task.
|
|
"""
|
|
CANCELLED_TASK_IDS.add(task_id)
|
|
cursor = conn.execute(
|
|
"""
|
|
UPDATE background_tasks
|
|
SET status = 'cancelled',
|
|
message = 'Задача отменена пользователем',
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ? AND user_id = ? AND status IN ('pending', 'running')
|
|
""",
|
|
(task_id, user_id),
|
|
)
|
|
return cursor.rowcount > 0
|
|
|
|
|
|
def is_task_cancelled(task_id: int, conn: Optional[sqlite3.Connection] = None) -> bool:
|
|
"""
|
|
Checks if a task was requested to be cancelled.
|
|
"""
|
|
if task_id in CANCELLED_TASK_IDS:
|
|
return True
|
|
if conn:
|
|
row = conn.execute(
|
|
"SELECT status FROM background_tasks WHERE id = ?", (task_id,)
|
|
).fetchone()
|
|
if row and row[0] == "cancelled":
|
|
CANCELLED_TASK_IDS.add(task_id)
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_task_status(conn: sqlite3.Connection, task_id: int, user_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Retrieves full details of a specific task.
|
|
"""
|
|
if user_id is not None:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT id, user_id, task_type, status, progress_current, progress_total,
|
|
message, params_json, result_json, error_message, created_at, updated_at
|
|
FROM background_tasks
|
|
WHERE id = ? AND user_id = ?
|
|
""",
|
|
(task_id, user_id),
|
|
).fetchone()
|
|
else:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT id, user_id, task_type, status, progress_current, progress_total,
|
|
message, params_json, result_json, error_message, created_at, updated_at
|
|
FROM background_tasks
|
|
WHERE id = ?
|
|
""",
|
|
(task_id,),
|
|
).fetchone()
|
|
|
|
if not row:
|
|
return None
|
|
|
|
return _format_task_row(row)
|
|
|
|
|
|
def get_active_task(conn: sqlite3.Connection, user_id: int, task_type: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Returns the most recent active (pending or running) task for user.
|
|
"""
|
|
if task_type:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT id, user_id, task_type, status, progress_current, progress_total,
|
|
message, params_json, result_json, error_message, created_at, updated_at
|
|
FROM background_tasks
|
|
WHERE user_id = ? AND status IN ('pending', 'running') AND task_type = ?
|
|
ORDER BY id DESC LIMIT 1
|
|
""",
|
|
(user_id, task_type),
|
|
).fetchone()
|
|
else:
|
|
row = conn.execute(
|
|
"""
|
|
SELECT id, user_id, task_type, status, progress_current, progress_total,
|
|
message, params_json, result_json, error_message, created_at, updated_at
|
|
FROM background_tasks
|
|
WHERE user_id = ? AND status IN ('pending', 'running')
|
|
ORDER BY id DESC LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
).fetchone()
|
|
|
|
if not row:
|
|
return None
|
|
|
|
return _format_task_row(row)
|
|
|
|
|
|
def list_user_tasks(conn: sqlite3.Connection, user_id: int, limit: int = 10) -> List[Dict[str, Any]]:
|
|
"""
|
|
Returns recent background tasks for user.
|
|
"""
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT id, user_id, task_type, status, progress_current, progress_total,
|
|
message, params_json, result_json, error_message, created_at, updated_at
|
|
FROM background_tasks
|
|
WHERE user_id = ?
|
|
ORDER BY id DESC LIMIT ?
|
|
""",
|
|
(user_id, limit),
|
|
).fetchall()
|
|
|
|
return [_format_task_row(r) for r in rows]
|
|
|
|
|
|
def _format_task_row(row: Any) -> Dict[str, Any]:
|
|
"""Helper to convert sqlite3.Row or tuple to a clean JSON-friendly dict."""
|
|
params = {}
|
|
result = {}
|
|
try:
|
|
params = json.loads(row["params_json"] or "{}")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
result = json.loads(row["result_json"] or "{}")
|
|
except Exception:
|
|
pass
|
|
|
|
total = row["progress_total"] or 0
|
|
current = row["progress_current"] or 0
|
|
percent = round((current / total) * 100) if total > 0 else 0
|
|
|
|
return {
|
|
"id": row["id"],
|
|
"user_id": row["user_id"],
|
|
"task_type": row["task_type"],
|
|
"status": row["status"],
|
|
"progress_current": current,
|
|
"progress_total": total,
|
|
"progress_percent": percent,
|
|
"message": row["message"] or "",
|
|
"params": params,
|
|
"result": result,
|
|
"error_message": row["error_message"] or "",
|
|
"created_at": str(row["created_at"]),
|
|
"updated_at": str(row["updated_at"]),
|
|
}
|