lou
224 lines
8.3 KiB
Python
224 lines
8.3 KiB
Python
"""
|
|
任务状态模块 (state)
|
|
模块:客户端 / 任务状态
|
|
输入:任务事件 (创建/发送成功/失败)
|
|
输出:SQLite 任务记录
|
|
|
|
SQLite 本地库, 重跑时查未完成任务 + 服务端 /chunks 对比出缺卷, 续传不重传。
|
|
|
|
表结构:
|
|
tasks task_id, file_name, file_size, chunk_count, status, created_at, updated_at
|
|
chunks task_id, idx, status, updated_at (只记录非 pending 状态, 缺行即未传)
|
|
"""
|
|
|
|
import argparse
|
|
import sqlite3
|
|
import sys
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
DEFAULT_DB = Path(__file__).resolve().parent / "config" / "tasks.db"
|
|
|
|
STATUS_UPLOADING = "uploading"
|
|
STATUS_DONE = "done"
|
|
STATUS_FAILED = "failed"
|
|
CHUNK_UPLOADED = "uploaded"
|
|
CHUNK_FAILED = "failed"
|
|
|
|
|
|
class TaskStore:
|
|
"""SQLite 任务状态存储 (每次操作独立连接, 文件锁安全)"""
|
|
|
|
def __init__(self, db_path: str | Path = DEFAULT_DB) -> None:
|
|
self.db_path = Path(db_path)
|
|
self._init_schema()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(self.db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
def _init_schema(self) -> None:
|
|
with self._connect() as conn:
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS tasks (
|
|
task_id TEXT PRIMARY KEY,
|
|
file_name TEXT NOT NULL,
|
|
file_size INTEGER NOT NULL,
|
|
chunk_count INTEGER NOT NULL,
|
|
status TEXT NOT NULL DEFAULT 'uploading',
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS chunks (
|
|
task_id TEXT NOT NULL,
|
|
idx INTEGER NOT NULL,
|
|
status TEXT NOT NULL DEFAULT 'uploaded',
|
|
updated_at TEXT NOT NULL,
|
|
PRIMARY KEY (task_id, idx)
|
|
);
|
|
"""
|
|
)
|
|
|
|
@staticmethod
|
|
def _now() -> str:
|
|
return datetime.now().isoformat(timespec="seconds")
|
|
|
|
# ---------- 任务事件 ----------
|
|
|
|
def create_task(
|
|
self, task_id: str, file_name: str, file_size: int, chunk_count: int
|
|
) -> None:
|
|
"""事件: 创建任务 (状态 uploading)"""
|
|
now = self._now()
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"INSERT INTO tasks (task_id, file_name, file_size, chunk_count, status, created_at, updated_at) "
|
|
"VALUES (?, ?, ?, ?, ?, ?, ?)",
|
|
(task_id, file_name, file_size, chunk_count, STATUS_UPLOADING, now, now),
|
|
)
|
|
|
|
def mark_chunk_uploaded(self, task_id: str, idx: int) -> None:
|
|
"""事件: 卷发送成功"""
|
|
self._upsert_chunk(task_id, idx, CHUNK_UPLOADED)
|
|
|
|
def mark_chunk_failed(self, task_id: str, idx: int) -> None:
|
|
"""事件: 卷发送失败 (重试耗尽, 待补传)"""
|
|
self._upsert_chunk(task_id, idx, CHUNK_FAILED)
|
|
|
|
def _upsert_chunk(self, task_id: str, idx: int, status: str) -> None:
|
|
now = self._now()
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"INSERT INTO chunks (task_id, idx, status, updated_at) VALUES (?, ?, ?, ?) "
|
|
"ON CONFLICT(task_id, idx) DO UPDATE SET status = excluded.status, updated_at = excluded.updated_at",
|
|
(task_id, idx, status, now),
|
|
)
|
|
conn.execute(
|
|
"UPDATE tasks SET updated_at = ? WHERE task_id = ?", (now, task_id)
|
|
)
|
|
|
|
def mark_done(self, task_id: str) -> None:
|
|
"""事件: 全部完成"""
|
|
self._set_task_status(task_id, STATUS_DONE)
|
|
|
|
def mark_failed(self, task_id: str) -> None:
|
|
"""事件: 任务失败"""
|
|
self._set_task_status(task_id, STATUS_FAILED)
|
|
|
|
def _set_task_status(self, task_id: str, status: str) -> None:
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"UPDATE tasks SET status = ?, updated_at = ? WHERE task_id = ?",
|
|
(status, self._now(), task_id),
|
|
)
|
|
|
|
# ---------- 查询 ----------
|
|
|
|
def get_task(self, task_id: str) -> dict[str, Any] | None:
|
|
with self._connect() as conn:
|
|
row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_chunk_status(self, task_id: str) -> dict[int, str]:
|
|
"""已记录的非 pending 卷状态 {idx: uploaded/failed}"""
|
|
with self._connect() as conn:
|
|
rows = conn.execute(
|
|
"SELECT idx, status FROM chunks WHERE task_id = ?", (task_id,)
|
|
).fetchall()
|
|
return {r["idx"]: r["status"] for r in rows}
|
|
|
|
def get_pending_tasks(self) -> list[dict[str, Any]]:
|
|
"""查询未完成任务 (重跑续传): 返回任务 + 已传卷集合"""
|
|
with self._connect() as conn:
|
|
rows = conn.execute(
|
|
"SELECT * FROM tasks WHERE status = ? ORDER BY created_at", (STATUS_UPLOADING,)
|
|
).fetchall()
|
|
tasks = [dict(r) for r in rows]
|
|
for t in tasks:
|
|
t["uploaded_chunks"] = sorted(
|
|
idx for idx, st in self.get_chunk_status(t["task_id"]).items()
|
|
if st == CHUNK_UPLOADED
|
|
)
|
|
return tasks
|
|
|
|
# ---------- 清理 ----------
|
|
|
|
def cleanup(self, days: int = 7) -> int:
|
|
"""清理超过 days 天未更新的非 uploading 任务, 返回删除数"""
|
|
cutoff = (datetime.now() - timedelta(days=days)).isoformat(timespec="seconds")
|
|
with self._connect() as conn:
|
|
rows = conn.execute(
|
|
"SELECT task_id FROM tasks WHERE status != ? AND updated_at < ?",
|
|
(STATUS_UPLOADING, cutoff),
|
|
).fetchall()
|
|
ids = [r["task_id"] for r in rows]
|
|
for tid in ids:
|
|
conn.execute("DELETE FROM chunks WHERE task_id = ?", (tid,))
|
|
conn.execute("DELETE FROM tasks WHERE task_id = ?", (tid,))
|
|
return len(ids)
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description="任务状态: SQLite 本地记录 (创建/卷状态/查询/清理)")
|
|
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
|
|
p_init = sub.add_parser("init", help="创建任务")
|
|
p_init.add_argument("task_id")
|
|
p_init.add_argument("file")
|
|
p_init.add_argument("--size", type=int, required=True)
|
|
p_init.add_argument("--chunks", type=int, required=True)
|
|
|
|
p_chunk = sub.add_parser("chunk", help="记录卷状态")
|
|
p_chunk.add_argument("task_id")
|
|
p_chunk.add_argument("index", type=int)
|
|
p_chunk.add_argument("--status", choices=["uploaded", "failed"], default="uploaded")
|
|
|
|
p_done = sub.add_parser("done", help="标记完成")
|
|
p_done.add_argument("task_id")
|
|
|
|
p_fail = sub.add_parser("failed", help="标记失败")
|
|
p_fail.add_argument("task_id")
|
|
|
|
sub.add_parser("pending", help="查询未完成任务")
|
|
|
|
p_clean = sub.add_parser("cleanup", help="清理旧任务")
|
|
p_clean.add_argument("--days", type=int, default=7)
|
|
|
|
args = ap.parse_args()
|
|
store = TaskStore()
|
|
|
|
if args.cmd == "init":
|
|
store.create_task(args.task_id, args.file, args.size, args.chunks)
|
|
print(f"[state] 任务 {args.task_id} 已创建 ({args.chunks} 卷)")
|
|
elif args.cmd == "chunk":
|
|
if args.status == "uploaded":
|
|
store.mark_chunk_uploaded(args.task_id, args.index)
|
|
else:
|
|
store.mark_chunk_failed(args.task_id, args.index)
|
|
print(f"[state] 卷 {args.index} -> {args.status}")
|
|
elif args.cmd == "done":
|
|
store.mark_done(args.task_id)
|
|
print(f"[state] 任务 {args.task_id} -> done")
|
|
elif args.cmd == "failed":
|
|
store.mark_failed(args.task_id)
|
|
print(f"[state] 任务 {args.task_id} -> failed")
|
|
elif args.cmd == "pending":
|
|
for t in store.get_pending_tasks():
|
|
print(
|
|
f" {t['task_id']} {t['file_name']} "
|
|
f"{len(t['uploaded_chunks'])}/{t['chunk_count']} 卷 "
|
|
f"(创建于 {t['created_at']})"
|
|
)
|
|
elif args.cmd == "cleanup":
|
|
n = store.cleanup(args.days)
|
|
print(f"[state] 清理 {n} 个过期任务")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|