aria2c for downloading

App/Youtube/youtubeRouter.py

@@ -2,10 +2,7 @@ from fastapi import APIRouter, HTTPException, BackgroundTasks
 import tempfile
 import os
 import logging
-import aiohttp
 import asyncio
-import aiofiles
-import subprocess
 from .Schema import YouTubeUploadTask
 
 # Configure logging
@@ -16,81 +13,48 @@ router = APIRouter(prefix="/youtube", tags=["youtube"])
 
 
 async def download_file(url: str) -> str:
-    """Download file from URL and save to temp directory with correct extension based on metadata."""
+    """Download file from URL using aria2c."""
     try:
-        async with aiohttp.ClientSession() as session:
-            async with session.get(url) as response:
-                if response.status not in (200, 206):
-                    raise HTTPException(
-                        status_code=400,
-                        detail=f"Failed to download file: HTTP {response.status}",
-                    )
-
-                # Create temp file with .tmp extension
-                fd, temp_path = tempfile.mkstemp(suffix=".tmp")
-
-                # Write content to temp file
-                async with aiofiles.open(fd, "wb") as tmp:
-                    while chunk := await response.content.read(8192):
-                        await tmp.write(chunk)
-
-                # Determine file format using ffprobe
-                try:
-                    # Get codec information
-                    codec_cmd = [
-                        "ffprobe",
-                        "-v",
-                        "error",
-                        "-select_streams",
-                        "v:0",
-                        "-show_entries",
-                        "stream=codec_name",
-                        "-of",
-                        "default=noprint_wrappers=1:nokey=1",
-                        temp_path,
-                    ]
-                    codec_result = subprocess.run(
-                        codec_cmd, capture_output=True, text=True
-                    )
-                    codec = codec_result.stdout.strip().lower()
-
-                    # Get container format
-                    format_cmd = [
-                        "ffprobe",
-                        "-v",
-                        "error",
-                        "-show_entries",
-                        "format=format_name",
-                        "-of",
-                        "default=noprint_wrappers=1:nokey=1",
-                        temp_path,
-                    ]
-                    format_result = subprocess.run(
-                        format_cmd, capture_output=True, text=True
-                    )
-                    format_name = format_result.stdout.strip().lower()
-
-                    # Determine extension based on codec and container
-                    if "mp4" in format_name or codec in ("h264", "aac"):
-                        ext = ".mp4"
-                    elif "webm" in format_name or codec == "vp9":
-                        ext = ".webm"
-                    elif "avi" in format_name:
-                        ext = ".avi"
-                    else:
-                        # Fallback to content-type if metadata unavailable
-                        content_type = response.headers.get("content-type", "")
-                        ext = ".mp4" if "video" in content_type else ".tmp"
-                except Exception:
-                    # Fallback if ffprobe fails
-                    content_type = response.headers.get("content-type", "")
-                    ext = ".mp4" if "video" in content_type else ".tmp"
-
-                # Rename file with correct extension
-                new_temp_path = f"{temp_path[:-4]}{ext}"
-                os.rename(temp_path, new_temp_path)
-
-                return new_temp_path
+        # Create temp directory for download
+        temp_dir = tempfile.mkdtemp()
+        temp_path = os.path.join(
+            temp_dir, "video"
+        )  # aria2c will add extension automatically
+
+        # Build aria2c command
+        command = [
+            "aria2c",
+            "--allow-overwrite=true",
+            "--auto-file-renaming=false",
+            "--max-connection-per-server=16",
+            "--split=16",
+            "--dir",
+            temp_dir,
+            "--out",
+            "video",
+            url,
+        ]
+
+        logger.info(f"Starting download with aria2c: {url}")
+        process = await asyncio.create_subprocess_exec(
+            *command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await process.communicate()
+
+        if process.returncode != 0:
+            raise Exception(f"Download failed: {stderr.decode()}")
+
+        # Find the downloaded file (aria2c adds extension automatically)
+        downloaded_files = os.listdir(temp_dir)
+        if not downloaded_files:
+            raise Exception("No file downloaded")
+
+        downloaded_file = os.path.join(temp_dir, downloaded_files[0])
+        logger.info(f"Download completed: {downloaded_file}")
+        return downloaded_file
+
     except Exception as e:
         logger.error(f"Error downloading file: {str(e)}")
         raise HTTPException(status_code=500, detail=f"File download failed: {str(e)}")
@@ -98,17 +62,18 @@ async def download_file(url: str) -> str:
 
 async def upload_video_background(task: YouTubeUploadTask):
     """Background task to handle video upload."""
-    temp_file = None
+    temp_dir = None
     try:
         logger.info(f"Starting download for video: {task.filename}")
-        temp_file = await download_file(task.filename)
-        logger.info(f"Download complete. Saved to: {temp_file}")
+        downloaded_file = await download_file(task.filename)
+        temp_dir = os.path.dirname(downloaded_file)
+        logger.info(f"Download complete. Saved to: {downloaded_file}")
 
         # Build the command
         command = [
             "/srv/youtube/youtubeuploader",
             "-filename",
-            temp_file,
+            downloaded_file,
             "-title",
             task.title,
             "-description",
@@ -143,12 +108,15 @@ async def upload_video_background(task: YouTubeUploadTask):
         logger.error(f"Error in upload process: {str(e)}")
         raise
     finally:
-        if temp_file and os.path.exists(temp_file):
+        # Clean up temp directory and its contents
+        if temp_dir and os.path.exists(temp_dir):
             try:
-                os.remove(temp_file)
-                logger.info(f"Cleaned up temporary file: {temp_file}")
+                for file in os.listdir(temp_dir):
+                    os.remove(os.path.join(temp_dir, file))
+                os.rmdir(temp_dir)
+                logger.info(f"Cleaned up temporary directory: {temp_dir}")
             except Exception as e:
-                logger.error(f"Failed to clean up temp file: {str(e)}")
+                logger.error(f"Failed to clean up temp directory: {str(e)}")
 
 
 @router.post("/upload")