Spaces:
Running
Running
File size: 2,328 Bytes
faefb1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | from __future__ import annotations
import asyncio
import logging
from pathlib import Path
from typing import List, Optional, Tuple
from urllib.parse import unquote, urlparse
import aiohttp
from app.config import get_settings
_logger = logging.getLogger(__name__)
_settings = get_settings()
def create_session_timeout(total_seconds: Optional[float] = None) -> aiohttp.ClientTimeout:
"""Create an aiohttp timeout using the configured request timeout."""
if total_seconds is None:
total_seconds = _settings.request_timeout_ms / 1000
return aiohttp.ClientTimeout(total=total_seconds)
async def download_url(
url: str,
*,
timeout_seconds: float = 120,
max_size_bytes: Optional[int] = None,
chunk_size: int = 256 * 1024,
) -> Tuple[bytes, Optional[str]]:
"""Download a file from a URL with streaming and size guard.
Returns (data, filename) where filename may be None.
Consolidates _download_file from csv_analysis_service and
_download_from_url from dataset_metadata_service.
"""
timeout = aiohttp.ClientTimeout(total=timeout_seconds)
try:
async with aiohttp.ClientSession(timeout=timeout) as session:
async with session.get(url) as resp:
if resp.status != 200:
raise RuntimeError(f"HTTP {resp.status} when fetching {url}")
if max_size_bytes and resp.content_length and resp.content_length > max_size_bytes:
raise RuntimeError(
f"Remote file advertises {resp.content_length} bytes, "
f"limit is {max_size_bytes}"
)
chunks: List[bytes] = []
total = 0
async for chunk in resp.content.iter_chunked(chunk_size):
total += len(chunk)
if max_size_bytes and total > max_size_bytes:
raise RuntimeError(f"Download exceeded {max_size_bytes} bytes")
chunks.append(chunk)
data = b"".join(chunks)
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
raise RuntimeError(f"Download failed for {url}: {exc}") from exc
parsed = urlparse(url)
filename = unquote(Path(parsed.path).name) if parsed.path else None
return data, filename
|