feat: B2 uit voor veilige archive-downloads
This commit is contained in:
Binary file not shown.
@@ -6,8 +6,8 @@ from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
VALID_STATUSES = {"queued", "running", "completed", "failed"}
|
||||
VALID_OPERATIONS = {"copy", "move"}
|
||||
VALID_STATUSES = {"queued", "running", "completed", "failed", "requested", "preparing", "ready"}
|
||||
VALID_OPERATIONS = {"copy", "move", "download"}
|
||||
TASK_MIGRATION_COLUMNS: dict[str, str] = {
|
||||
"operation": "TEXT NOT NULL DEFAULT 'copy'",
|
||||
"status": "TEXT NOT NULL DEFAULT 'queued'",
|
||||
@@ -32,9 +32,18 @@ class TaskRepository:
|
||||
self._db_path = db_path
|
||||
self._ensure_schema()
|
||||
|
||||
def create_task(self, operation: str, source: str, destination: str, task_id: str | None = None) -> dict:
|
||||
def create_task(
|
||||
self,
|
||||
operation: str,
|
||||
source: str,
|
||||
destination: str,
|
||||
task_id: str | None = None,
|
||||
status: str = "queued",
|
||||
) -> dict:
|
||||
if operation not in VALID_OPERATIONS:
|
||||
raise ValueError("invalid operation")
|
||||
if status not in VALID_STATUSES:
|
||||
raise ValueError("invalid status")
|
||||
|
||||
task_id = task_id or str(uuid.uuid4())
|
||||
created_at = self._now_iso()
|
||||
@@ -52,7 +61,7 @@ class TaskRepository:
|
||||
(
|
||||
task_id,
|
||||
operation,
|
||||
"queued",
|
||||
status,
|
||||
source,
|
||||
destination,
|
||||
None,
|
||||
@@ -145,6 +154,24 @@ class TaskRepository:
|
||||
("running", started_at, done_bytes, total_bytes, done_items, total_items, current_item, task_id),
|
||||
)
|
||||
|
||||
def mark_preparing(
|
||||
self,
|
||||
task_id: str,
|
||||
done_items: int | None = None,
|
||||
total_items: int | None = None,
|
||||
current_item: str | None = None,
|
||||
) -> None:
|
||||
started_at = self._now_iso()
|
||||
with self._connection() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE tasks
|
||||
SET status = ?, started_at = COALESCE(started_at, ?), done_items = ?, total_items = ?, current_item = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
("preparing", started_at, done_items, total_items, current_item, task_id),
|
||||
)
|
||||
|
||||
def update_progress(
|
||||
self,
|
||||
task_id: str,
|
||||
@@ -183,6 +210,23 @@ class TaskRepository:
|
||||
("completed", finished_at, done_bytes, total_bytes, done_items, total_items, task_id),
|
||||
)
|
||||
|
||||
def mark_ready(
|
||||
self,
|
||||
task_id: str,
|
||||
done_items: int | None = None,
|
||||
total_items: int | None = None,
|
||||
) -> None:
|
||||
finished_at = self._now_iso()
|
||||
with self._connection() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE tasks
|
||||
SET status = ?, finished_at = ?, done_items = ?, total_items = ?, current_item = NULL
|
||||
WHERE id = ?
|
||||
""",
|
||||
("ready", finished_at, done_items, total_items, task_id),
|
||||
)
|
||||
|
||||
def mark_failed(
|
||||
self,
|
||||
task_id: str,
|
||||
@@ -244,14 +288,62 @@ class TaskRepository:
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS task_artifacts (
|
||||
task_id TEXT PRIMARY KEY,
|
||||
file_path TEXT NOT NULL,
|
||||
file_name TEXT NOT NULL,
|
||||
expires_at TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_created_at_desc
|
||||
ON tasks(created_at DESC)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS idx_task_artifacts_expires_at
|
||||
ON task_artifacts(expires_at ASC)
|
||||
"""
|
||||
)
|
||||
self._migrate_tasks_columns(conn)
|
||||
|
||||
def upsert_artifact(self, *, task_id: str, file_path: str, file_name: str, expires_at: str) -> dict:
|
||||
created_at = self._now_iso()
|
||||
with self._connection() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO task_artifacts (task_id, file_path, file_name, expires_at, created_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(task_id) DO UPDATE SET
|
||||
file_path = excluded.file_path,
|
||||
file_name = excluded.file_name,
|
||||
expires_at = excluded.expires_at
|
||||
""",
|
||||
(task_id, file_path, file_name, expires_at, created_at),
|
||||
)
|
||||
row = conn.execute("SELECT * FROM task_artifacts WHERE task_id = ?", (task_id,)).fetchone()
|
||||
return self._artifact_to_dict(row)
|
||||
|
||||
def get_artifact(self, task_id: str) -> dict | None:
|
||||
with self._connection() as conn:
|
||||
row = conn.execute("SELECT * FROM task_artifacts WHERE task_id = ?", (task_id,)).fetchone()
|
||||
return self._artifact_to_dict(row) if row else None
|
||||
|
||||
def list_artifacts(self) -> list[dict]:
|
||||
with self._connection() as conn:
|
||||
rows = conn.execute("SELECT * FROM task_artifacts ORDER BY created_at ASC").fetchall()
|
||||
return [self._artifact_to_dict(row) for row in rows]
|
||||
|
||||
def delete_artifact(self, task_id: str) -> None:
|
||||
with self._connection() as conn:
|
||||
conn.execute("DELETE FROM task_artifacts WHERE task_id = ?", (task_id,))
|
||||
|
||||
def _migrate_tasks_columns(self, conn: sqlite3.Connection) -> None:
|
||||
rows = conn.execute("PRAGMA table_info(tasks)").fetchall()
|
||||
existing_columns = {row["name"] for row in rows}
|
||||
@@ -298,6 +390,16 @@ class TaskRepository:
|
||||
"finished_at": row["finished_at"],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _artifact_to_dict(row: sqlite3.Row) -> dict:
|
||||
return {
|
||||
"task_id": row["task_id"],
|
||||
"file_path": row["file_path"],
|
||||
"file_name": row["file_name"],
|
||||
"expires_at": row["expires_at"],
|
||||
"created_at": row["created_at"],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(tz=timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
Reference in New Issue
Block a user