refactor: unify resumable collector downloads

backend/app/services/collectors/iptoasn.py

@@ -8,9 +8,6 @@ from __future__ import annotations
 
 import asyncio
 import gzip
-import hashlib
-import json
-import tempfile
 import time
 from datetime import UTC, datetime
 from ipaddress import summarize_address_range, ip_address
@@ -20,6 +17,7 @@ from typing import Any
 import httpx
 
 from app.services.collectors.base import BaseCollector
+from app.services.collectors.downloads import ResumableFileDownloader
 
 
 class IPtoASNPrefixGeoCollector(BaseCollector):
@@ -29,7 +27,10 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
     frequency_hours = 24
     data_type = "prefix_geography"
     fail_on_empty = True
-    _cache_dir = Path(tempfile.gettempdir()) / "planet-download-cache" / "iptoasn"
+    _downloader = ResumableFileDownloader(
+        cache_namespace="iptoasn",
+        default_accept="application/gzip,application/octet-stream,*/*",
+    )
 
     @staticmethod
     def _build_dataset_urls(resolved_url: str) -> list[str]:
@@ -40,57 +41,6 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
             ]
         return [resolved_url]
 
-    @staticmethod
-    def _cache_key(url: str) -> str:
-        return hashlib.sha1(url.encode("utf-8")).hexdigest()[:16]
-
-    @classmethod
-    def _cache_paths(cls, url: str) -> tuple[Path, Path, Path]:
-        key = cls._cache_key(url)
-        gz_path = cls._cache_dir / f"{key}.tsv.gz"
-        part_path = cls._cache_dir / f"{key}.tsv.gz.part"
-        meta_path = cls._cache_dir / f"{key}.meta.json"
-        return gz_path, part_path, meta_path
-
-    @staticmethod
-    def _load_meta(meta_path: Path) -> dict[str, Any]:
-        if not meta_path.exists():
-            return {}
-        try:
-            return json.loads(meta_path.read_text(encoding="utf-8"))
-        except (json.JSONDecodeError, OSError):
-            return {}
-
-    @staticmethod
-    def _save_meta(meta_path: Path, payload: dict[str, Any]) -> None:
-        meta_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
-
-    @staticmethod
-    def _validators_match(meta: dict[str, Any], remote: dict[str, Any]) -> bool:
-        etag = str(remote.get("etag") or "").strip()
-        last_modified = str(remote.get("last_modified") or "").strip()
-        if etag:
-            return etag == str(meta.get("etag") or "").strip()
-        if last_modified:
-            return last_modified == str(meta.get("last_modified") or "").strip()
-        return True
-
-    async def _fetch_remote_info(self, client: httpx.AsyncClient, url: str) -> dict[str, Any]:
-        try:
-            response = await client.head(url)
-            if response.status_code >= 400:
-                return {}
-            content_length_raw = response.headers.get("content-length")
-            content_length = int(content_length_raw) if content_length_raw else None
-            return {
-                "etag": response.headers.get("etag"),
-                "last_modified": response.headers.get("last-modified"),
-                "content_length": content_length,
-                "accept_ranges": (response.headers.get("accept-ranges") or "").lower(),
-            }
-        except (httpx.HTTPError, ValueError):
-            return {}
-
     def _parse_rows_from_gzip_file(self, file_path: Path) -> list[dict[str, Any]]:
         rows: list[dict[str, Any]] = []
         with gzip.open(file_path, "rt", encoding="utf-8", errors="replace") as f:
@@ -115,94 +65,6 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
             raise RuntimeError(f"IPtoASN dataset parsed empty rows: {file_path.name}")
         return rows
 
-    async def _download_dataset_with_resume(
-        self,
-        client: httpx.AsyncClient,
-        url: str,
-        *,
-        progress_callback=None,
-    ) -> Path:
-        self._cache_dir.mkdir(parents=True, exist_ok=True)
-        gz_path, part_path, meta_path = self._cache_paths(url)
-        meta = self._load_meta(meta_path)
-        remote = await self._fetch_remote_info(client, url)
-
-        # Reuse existing complete file when validators and size still match.
-        if gz_path.exists():
-            local_size = gz_path.stat().st_size
-            remote_size = remote.get("content_length")
-            if (
-                self._validators_match(meta, remote)
-                and (remote_size is None or local_size == remote_size)
-            ):
-                try:
-                    self._parse_rows_from_gzip_file(gz_path)
-                    if progress_callback and local_size > 0:
-                        await progress_callback(local_size)
-                    return gz_path
-                except Exception:
-                    gz_path.unlink(missing_ok=True)
-
-        expected_size = remote.get("content_length")
-        can_resume = (remote.get("accept_ranges") or "") == "bytes"
-        resume_from = part_path.stat().st_size if part_path.exists() else 0
-        if expected_size is not None and resume_from > expected_size:
-            part_path.unlink(missing_ok=True)
-            resume_from = 0
-        if not self._validators_match(meta, remote):
-            part_path.unlink(missing_ok=True)
-            resume_from = 0
-
-        headers = {
-            "User-Agent": "Planet-Intelligence-System/1.0 (Python/collector)",
-            "Accept": "application/gzip,application/octet-stream,*/*",
-        }
-        if can_resume and resume_from > 0:
-            headers["Range"] = f"bytes={resume_from}-"
-            if remote.get("etag"):
-                headers["If-Range"] = str(remote.get("etag"))
-            elif remote.get("last_modified"):
-                headers["If-Range"] = str(remote.get("last_modified"))
-
-        async with client.stream("GET", url, headers=headers) as response:
-            response.raise_for_status()
-
-            if response.status_code == 206 and resume_from > 0:
-                mode = "ab"
-            else:
-                mode = "wb"
-                resume_from = 0
-
-            with part_path.open(mode) as f:
-                downloaded = resume_from
-                if progress_callback and downloaded > 0:
-                    await progress_callback(downloaded)
-                async for chunk in response.aiter_bytes():
-                    if chunk:
-                        f.write(chunk)
-                        downloaded += len(chunk)
-                        if progress_callback:
-                            await progress_callback(downloaded)
-
-        part_size = part_path.stat().st_size if part_path.exists() else 0
-        if expected_size is not None and part_size != expected_size:
-            raise RuntimeError(
-                f"IPtoASN download incomplete for {url}: expected={expected_size}, got={part_size}"
-            )
-
-        part_path.replace(gz_path)
-        self._save_meta(
-            meta_path,
-            {
-                "url": url,
-                "etag": remote.get("etag"),
-                "last_modified": remote.get("last_modified"),
-                "content_length": expected_size,
-                "updated_at": datetime.now(UTC).isoformat(),
-            },
-        )
-        return gz_path
-
     async def _fetch_dataset_rows(
         self,
         client: httpx.AsyncClient,
@@ -210,27 +72,21 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
         *,
         progress_callback=None,
     ) -> list[dict[str, Any]]:
-        try:
-            gz_path = await self._download_dataset_with_resume(
-                client,
-                url,
-                progress_callback=progress_callback,
-            )
-            return self._parse_rows_from_gzip_file(gz_path)
-        except Exception:
-            # Bad partial/complete cache or interrupted stream: clear local files
-            # and retry once from scratch.
-            gz_path, part_path, meta_path = self._cache_paths(url)
-            gz_path.unlink(missing_ok=True)
-            part_path.unlink(missing_ok=True)
-            meta_path.unlink(missing_ok=True)
+        file_path = await self._downloader.download_file(
+            client,
+            url,
+            extension=".tsv.gz",
+            progress_callback=progress_callback,
+            validate_existing=lambda p: self._validate_gzip_dataset(p),
+        )
+        return self._parse_rows_from_gzip_file(file_path)
 
-            gz_path = await self._download_dataset_with_resume(
-                client,
-                url,
-                progress_callback=progress_callback,
-            )
-            return self._parse_rows_from_gzip_file(gz_path)
+    def _validate_gzip_dataset(self, file_path: Path) -> bool:
+        try:
+            self._parse_rows_from_gzip_file(file_path)
+            return True
+        except Exception:
+            return False
 
     async def fetch(self) -> list[dict[str, Any]]:
         if not self._resolved_url:
@@ -240,7 +96,7 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
 
         async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client:
             remote_infos = await asyncio.gather(
-                *(self._fetch_remote_info(client, url) for url in dataset_urls)
+                *(self._downloader.fetch_remote_info(client, url) for url in dataset_urls)
             )
             expected_sizes = [
                 info.get("content_length")
@@ -260,7 +116,7 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
             last_emit = {"t": 0.0, "value": 0}
             min_emit_bytes = max(total_expected // 200, 2 * 1024 * 1024) if total_expected > 0 else 4 * 1024 * 1024
 
-            async def on_url_progress(url: str, downloaded_bytes: int) -> None:
+            async def on_url_progress(url: str, downloaded_bytes: int, total_bytes: int | None) -> None:
                 if total_expected <= 0:
                     return
                 async with progress_lock:
@@ -286,7 +142,7 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
                     self._fetch_dataset_rows(
                         client,
                         url,
-                        progress_callback=lambda downloaded, u=url: on_url_progress(u, downloaded),
+                        progress_callback=lambda downloaded, total, u=url: on_url_progress(u, downloaded, total),
                     )
                     for url in dataset_urls
                 )