91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
"""Download endpoint for split results."""
|
|
|
|
from fastapi import APIRouter, HTTPException
|
|
from fastapi.responses import FileResponse
|
|
from zipfile import ZipFile
|
|
import os
|
|
|
|
from backend.config import settings
|
|
from backend.services.task_manager import task_manager
|
|
from backend.models.response import TaskStatus
|
|
|
|
router = APIRouter(prefix="/api", tags=["download"])
|
|
|
|
|
|
def _prepare_download(task_id: str):
|
|
"""Common logic to prepare and return the ZIP file."""
|
|
if not task_manager.has_task(task_id):
|
|
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
|
|
|
|
status = task_manager.get_status(task_id)
|
|
if status["status"] != TaskStatus.DONE:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Task {task_id} is not complete. Current status: {status['status']}"
|
|
)
|
|
|
|
output_dir = settings.temp_dir / task_id / "output"
|
|
if not output_dir.exists() or not any(output_dir.iterdir()):
|
|
raise HTTPException(status_code=404, detail="No output files found")
|
|
|
|
zip_path = settings.temp_dir / task_id / "splits.zip"
|
|
with ZipFile(zip_path, "w") as zipf:
|
|
for file_path in output_dir.iterdir():
|
|
if file_path.is_file():
|
|
zipf.write(file_path, arcname=file_path.name)
|
|
|
|
return zip_path
|
|
|
|
|
|
@router.get("/download/{task_id}")
|
|
async def download_results(task_id: str):
|
|
zip_path = _prepare_download(task_id)
|
|
return FileResponse(
|
|
path=zip_path,
|
|
media_type="application/zip",
|
|
filename=f"{task_id}_splits.zip",
|
|
headers={"Content-Disposition": f"attachment; filename={task_id}_splits.zip"}
|
|
)
|
|
|
|
|
|
@router.get("/download/{task_id}/splits.zip")
|
|
async def download_results_as_zip(task_id: str):
|
|
"""Alias endpoint that provides a .zip suffix for easier curl usage."""
|
|
zip_path = _prepare_download(task_id)
|
|
return FileResponse(
|
|
path=zip_path,
|
|
media_type="application/zip",
|
|
filename="splits.zip",
|
|
headers={"Content-Disposition": "attachment; filename=splits.zip"}
|
|
)
|
|
|
|
|
|
# NEW: Endpoint for downloading individual tracks
|
|
@router.get("/download/{task_id}/{filename}")
|
|
async def download_single_track(task_id: str, filename: str):
|
|
"""
|
|
Download a single track file from the output directory.
|
|
"""
|
|
if not task_manager.has_task(task_id):
|
|
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
|
|
|
|
status = task_manager.get_status(task_id)
|
|
if status["status"] != TaskStatus.DONE:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Task {task_id} is not complete. Current status: {status['status']}"
|
|
)
|
|
|
|
output_dir = settings.temp_dir / task_id / "output"
|
|
file_path = output_dir / filename
|
|
|
|
if not file_path.exists() or not file_path.is_file():
|
|
raise HTTPException(status_code=404, detail=f"File '{filename}' not found")
|
|
|
|
return FileResponse(
|
|
path=file_path,
|
|
filename=filename,
|
|
media_type="application/octet-stream",
|
|
headers={"Content-Disposition": f"attachment; filename={filename}"}
|
|
)
|