Spaces:
Running
Running
File size: 6,757 Bytes
faefb1f 722c296 faefb1f 722c296 faefb1f 722c296 faefb1f 722c296 faefb1f 722c296 faefb1f 722c296 bd469c1 722c296 bd469c1 722c296 faefb1f bd469c1 faefb1f bd469c1 faefb1f 722c296 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | from __future__ import annotations
import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple
from urllib.parse import unquote, urlparse
import aiohttp
import httpx
from app.config import get_settings
_logger = logging.getLogger(__name__)
_settings = get_settings()
class DownloadError(RuntimeError):
"""Raised when a remote download fails (HTTP error, timeout, or size limit).
``is_size_error`` distinguishes size-limit violations so callers can map
the failure to the appropriate domain exception.
"""
def __init__(self, message: str, *, is_size_error: bool = False) -> None:
super().__init__(message)
self.is_size_error = is_size_error
class SharedAsyncClient:
"""Thread-safe, lazily-created, connection-pooled ``httpx.AsyncClient``.
Several services (GCS, Google Maps, Google OAuth, scheduler) previously
duplicated the same double-checked-lock initialisation and shutdown logic.
This helper centralises that pattern so a single instance is created on
first use and safely closed exactly once on shutdown.
"""
def __init__(
self,
*,
timeout: Any = 30.0,
limits: Optional[httpx.Limits] = None,
follow_redirects: bool = True,
) -> None:
self._client: Optional[httpx.AsyncClient] = None
self._lock = asyncio.Lock()
self._timeout = timeout
self._limits = limits or httpx.Limits(
max_connections=100,
max_keepalive_connections=20,
keepalive_expiry=30,
)
self._follow_redirects = follow_redirects
async def get(self) -> httpx.AsyncClient:
"""Return the shared client, creating it lazily on first use."""
if self._client is None or self._client.is_closed:
async with self._lock:
if self._client is None or self._client.is_closed:
self._client = httpx.AsyncClient(
timeout=self._timeout,
follow_redirects=self._follow_redirects,
limits=self._limits,
)
return self._client
async def close(self) -> None:
"""Close the underlying client and release pooled connections."""
async with self._lock:
if self._client is not None and not self._client.is_closed:
await self._client.aclose()
self._client = None
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)
class SharedAiohttpSession:
"""Reuse a single ``aiohttp.ClientSession`` per event loop.
aiohttp sessions are bound to the loop they are created on, so this caches
one session per event loop. In the long-lived server loop this removes the
per-request session creation (and the associated TCP/TLS handshake) from hot
paths such as chat completions and web search.
"""
def __init__(self) -> None:
self._sessions: Dict[int, aiohttp.ClientSession] = {}
async def get(self, timeout: Optional[aiohttp.ClientTimeout] = None) -> aiohttp.ClientSession:
loop = asyncio.get_running_loop()
loop_id = id(loop)
session = self._sessions.get(loop_id)
if session is None or session.closed:
session = aiohttp.ClientSession(timeout=timeout or create_session_timeout())
self._sessions[loop_id] = session
return session
async def close_all(self) -> None:
sessions = list(self._sessions.values())
self._sessions.clear()
for session in sessions:
if session is not None and not session.closed:
await session.close()
_shared_aiohttp_session = SharedAiohttpSession()
async def get_shared_aiohttp_session(
timeout: Optional[aiohttp.ClientTimeout] = None,
) -> aiohttp.ClientSession:
"""Return the loop-bound shared aiohttp session, created lazily on first use."""
return await _shared_aiohttp_session.get(timeout)
@asynccontextmanager
async def shared_aiohttp_session(
timeout: Optional[aiohttp.ClientTimeout] = None,
) -> AsyncIterator[aiohttp.ClientSession]:
"""Context manager that yields the shared session but does NOT close it on exit.
Drop-in replacement for ``async with aiohttp.ClientSession(...)`` so call
sites keep the same structure while reusing one pooled session per loop.
"""
session = await get_shared_aiohttp_session(timeout)
yield session
async def close_shared_aiohttp_sessions() -> None:
"""Close all cached sessions. Call once on application shutdown."""
await _shared_aiohttp_session.close_all()
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:
session = await get_shared_aiohttp_session()
async with session.get(url, timeout=timeout) as resp:
if resp.status != 200:
raise DownloadError(f"HTTP {resp.status} when fetching {url}")
if max_size_bytes and resp.content_length and resp.content_length > max_size_bytes:
raise DownloadError(
f"Remote file advertises {resp.content_length} bytes, "
f"limit is {max_size_bytes}",
is_size_error=True,
)
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 DownloadError(
f"Download exceeded {max_size_bytes} bytes",
is_size_error=True,
)
chunks.append(chunk)
data = b"".join(chunks)
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
raise DownloadError(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
|