refactor: unify resumable collector downloads

This commit is contained in:
linkong
2026-04-07 16:11:34 +08:00
parent 31672b7ba2
commit 3f5505f03e
4 changed files with 268 additions and 342 deletions

View File

@@ -0,0 +1,204 @@
"""Shared resumable download helpers for collectors."""
from __future__ import annotations
import hashlib
import json
import tempfile
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Awaitable, Callable
import httpx
ProgressCallback = Callable[[int, int | None], Awaitable[None]]
ValidateCallback = Callable[[Path], bool]
class ResumableFileDownloader:
"""Download files with cache validators and byte-range resume support."""
def __init__(
self,
*,
cache_namespace: str,
user_agent: str = "Planet-Intelligence-System/1.0 (Python/collector)",
default_accept: str = "*/*",
) -> None:
self._cache_dir = Path(tempfile.gettempdir()) / "planet-download-cache" / cache_namespace
self._user_agent = user_agent
self._default_accept = default_accept
@staticmethod
def _cache_key(url: str) -> str:
return hashlib.sha1(url.encode("utf-8")).hexdigest()[:16]
def _cache_paths(self, url: str, extension: str) -> tuple[Path, Path, Path]:
key = self._cache_key(url)
normalized_ext = extension if extension.startswith(".") else f".{extension}"
final_path = self._cache_dir / f"{key}{normalized_ext}"
part_path = self._cache_dir / f"{key}{normalized_ext}.part"
meta_path = self._cache_dir / f"{key}.meta.json"
return final_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 {}
async def download_file(
self,
client: httpx.AsyncClient,
url: str,
*,
extension: str,
accept: str | None = None,
progress_callback: ProgressCallback | None = None,
validate_existing: ValidateCallback | None = None,
) -> Path:
self._cache_dir.mkdir(parents=True, exist_ok=True)
final_path, part_path, meta_path = self._cache_paths(url, extension)
meta = self._load_meta(meta_path)
remote = await self.fetch_remote_info(client, url)
expected_size = remote.get("content_length")
if final_path.exists():
local_size = final_path.stat().st_size
size_match = expected_size is None or local_size == expected_size
if self._validators_match(meta, remote) and size_match:
if validate_existing and not validate_existing(final_path):
final_path.unlink(missing_ok=True)
else:
if progress_callback and expected_size and expected_size > 0:
await progress_callback(expected_size, expected_size)
return final_path
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": self._user_agent,
"Accept": accept or self._default_accept,
}
if final_path.exists():
if meta.get("etag"):
headers["If-None-Match"] = str(meta.get("etag"))
elif meta.get("last_modified"):
headers["If-Modified-Since"] = str(meta.get("last_modified"))
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:
if response.status_code == 304 and final_path.exists():
if progress_callback and expected_size and expected_size > 0:
await progress_callback(expected_size, expected_size)
return final_path
response.raise_for_status()
if response.status_code == 206 and resume_from > 0:
mode = "ab"
else:
mode = "wb"
resume_from = 0
downloaded = resume_from
last_emit_bytes = 0
last_emit_time = time.monotonic()
min_emit_bytes = (
max(expected_size // 150, 512 * 1024) if expected_size and expected_size > 0 else 1024 * 1024
)
with part_path.open(mode) as f:
if progress_callback and downloaded > 0:
await progress_callback(downloaded, expected_size)
async for chunk in response.aiter_bytes():
if not chunk:
continue
f.write(chunk)
downloaded += len(chunk)
if not progress_callback:
continue
now = time.monotonic()
should_emit = (
expected_size is None
or downloaded >= expected_size
or downloaded - last_emit_bytes >= min_emit_bytes
or now - last_emit_time >= 2.0
)
if should_emit:
last_emit_bytes = downloaded
last_emit_time = now
await progress_callback(downloaded, expected_size)
final_size = part_path.stat().st_size if part_path.exists() else 0
if expected_size is not None and final_size != expected_size:
raise RuntimeError(
f"Resumable download incomplete for {url}: expected={expected_size}, got={final_size}"
)
part_path.replace(final_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(),
},
)
if validate_existing and not validate_existing(final_path):
raise RuntimeError(f"Downloaded file validation failed for {url}")
if progress_callback and expected_size and expected_size > 0:
await progress_callback(expected_size, expected_size)
return final_path

View File

@@ -8,9 +8,6 @@ from __future__ import annotations
import asyncio import asyncio
import gzip import gzip
import hashlib
import json
import tempfile
import time import time
from datetime import UTC, datetime from datetime import UTC, datetime
from ipaddress import summarize_address_range, ip_address from ipaddress import summarize_address_range, ip_address
@@ -20,6 +17,7 @@ from typing import Any
import httpx import httpx
from app.services.collectors.base import BaseCollector from app.services.collectors.base import BaseCollector
from app.services.collectors.downloads import ResumableFileDownloader
class IPtoASNPrefixGeoCollector(BaseCollector): class IPtoASNPrefixGeoCollector(BaseCollector):
@@ -29,7 +27,10 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
frequency_hours = 24 frequency_hours = 24
data_type = "prefix_geography" data_type = "prefix_geography"
fail_on_empty = True 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 @staticmethod
def _build_dataset_urls(resolved_url: str) -> list[str]: def _build_dataset_urls(resolved_url: str) -> list[str]:
@@ -40,57 +41,6 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
] ]
return [resolved_url] 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]]: def _parse_rows_from_gzip_file(self, file_path: Path) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []
with gzip.open(file_path, "rt", encoding="utf-8", errors="replace") as f: 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}") raise RuntimeError(f"IPtoASN dataset parsed empty rows: {file_path.name}")
return rows 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( async def _fetch_dataset_rows(
self, self,
client: httpx.AsyncClient, client: httpx.AsyncClient,
@@ -210,27 +72,21 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
*, *,
progress_callback=None, progress_callback=None,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
try: file_path = await self._downloader.download_file(
gz_path = await self._download_dataset_with_resume( client,
client, url,
url, extension=".tsv.gz",
progress_callback=progress_callback, progress_callback=progress_callback,
) validate_existing=lambda p: self._validate_gzip_dataset(p),
return self._parse_rows_from_gzip_file(gz_path) )
except Exception: return self._parse_rows_from_gzip_file(file_path)
# 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)
gz_path = await self._download_dataset_with_resume( def _validate_gzip_dataset(self, file_path: Path) -> bool:
client, try:
url, self._parse_rows_from_gzip_file(file_path)
progress_callback=progress_callback, return True
) except Exception:
return self._parse_rows_from_gzip_file(gz_path) return False
async def fetch(self) -> list[dict[str, Any]]: async def fetch(self) -> list[dict[str, Any]]:
if not self._resolved_url: if not self._resolved_url:
@@ -240,7 +96,7 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client: async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client:
remote_infos = await asyncio.gather( 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 = [ expected_sizes = [
info.get("content_length") info.get("content_length")
@@ -260,7 +116,7 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
last_emit = {"t": 0.0, "value": 0} 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 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: if total_expected <= 0:
return return
async with progress_lock: async with progress_lock:
@@ -286,7 +142,7 @@ class IPtoASNPrefixGeoCollector(BaseCollector):
self._fetch_dataset_rows( self._fetch_dataset_rows(
client, client,
url, 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 for url in dataset_urls
) )

View File

@@ -7,17 +7,13 @@ allocation geography as prefix-centric fallback hints.
from __future__ import annotations from __future__ import annotations
import ipaddress import ipaddress
import hashlib
import json
import tempfile
import time
from datetime import UTC, datetime from datetime import UTC, datetime
from pathlib import Path
from typing import Any from typing import Any
import httpx import httpx
from app.services.collectors.base import BaseCollector from app.services.collectors.base import BaseCollector
from app.services.collectors.downloads import ResumableFileDownloader
class NRODelegatedPrefixGeoCollector(BaseCollector): class NRODelegatedPrefixGeoCollector(BaseCollector):
@@ -27,168 +23,17 @@ class NRODelegatedPrefixGeoCollector(BaseCollector):
frequency_hours = 24 frequency_hours = 24
data_type = "prefix_geography" data_type = "prefix_geography"
fail_on_empty = True fail_on_empty = True
_cache_dir = Path(tempfile.gettempdir()) / "planet-download-cache" / "nro" _downloader = ResumableFileDownloader(
cache_namespace="nro",
@staticmethod default_accept="text/plain,*/*",
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)
txt_path = cls._cache_dir / f"{key}.txt"
part_path = cls._cache_dir / f"{key}.txt.part"
meta_path = cls._cache_dir / f"{key}.meta.json"
return txt_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 {}
async def _download_body_with_resume(self, client: httpx.AsyncClient, url: str) -> str:
self._cache_dir.mkdir(parents=True, exist_ok=True)
txt_path, part_path, meta_path = self._cache_paths(url)
meta = self._load_meta(meta_path)
remote = await self._fetch_remote_info(client, url)
if txt_path.exists():
local_size = txt_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):
if remote_size > 0:
await self.update_progress(remote_size, commit=True, force=True)
return txt_path.read_text(encoding="utf-8", errors="replace")
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": "text/plain,*/*",
}
if txt_path.exists():
if meta.get("etag"):
headers["If-None-Match"] = str(meta.get("etag"))
elif meta.get("last_modified"):
headers["If-Modified-Since"] = str(meta.get("last_modified"))
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:
if response.status_code == 304 and txt_path.exists():
if expected_size and expected_size > 0:
await self.update_progress(expected_size, commit=True, force=True)
return txt_path.read_text(encoding="utf-8", errors="replace")
response.raise_for_status()
if response.status_code == 206 and resume_from > 0:
mode = "ab"
else:
mode = "wb"
resume_from = 0
downloaded = resume_from
last_emit = 0
last_emit_time = time.monotonic()
min_emit_bytes = max(expected_size // 150, 512 * 1024) if expected_size and expected_size > 0 else 1024 * 1024
with part_path.open(mode) as f:
if downloaded > 0 and expected_size and expected_size > 0:
await self.update_progress(min(downloaded, expected_size), commit=True)
async for chunk in response.aiter_bytes():
if not chunk:
continue
f.write(chunk)
downloaded += len(chunk)
if not expected_size or expected_size <= 0:
continue
now = time.monotonic()
should_emit = (
downloaded >= expected_size
or downloaded - last_emit >= min_emit_bytes
or now - last_emit_time >= 2.0
)
if should_emit:
last_emit = downloaded
last_emit_time = now
await self.update_progress(min(downloaded, expected_size), commit=True)
final_size = part_path.stat().st_size if part_path.exists() else 0
if expected_size is not None and final_size != expected_size:
raise RuntimeError(
f"NRO download incomplete for {url}: expected={expected_size}, got={final_size}"
)
part_path.replace(txt_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(),
},
)
if expected_size and expected_size > 0:
await self.update_progress(expected_size, commit=True, force=True)
return txt_path.read_text(encoding="utf-8", errors="replace")
async def fetch(self) -> list[dict[str, Any]]: async def fetch(self) -> list[dict[str, Any]]:
if not self._resolved_url: if not self._resolved_url:
raise RuntimeError("NRO delegated stats URL is not configured") raise RuntimeError("NRO delegated stats URL is not configured")
async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client: async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client:
remote = await self._fetch_remote_info(client, self._resolved_url) remote = await self._downloader.fetch_remote_info(client, self._resolved_url)
total_expected = remote.get("content_length") or 0 total_expected = remote.get("content_length") or 0
if total_expected > 0 and self._current_task and self._db_session: if total_expected > 0 and self._current_task and self._db_session:
self._current_task.total_records = total_expected self._current_task.total_records = total_expected
@@ -197,14 +42,18 @@ class NRODelegatedPrefixGeoCollector(BaseCollector):
await self._db_session.commit() await self._db_session.commit()
await self._publish_task_update(force=True) await self._publish_task_update(force=True)
try: async def on_progress(downloaded: int, total: int | None) -> None:
body = await self._download_body_with_resume(client, self._resolved_url) if not total or total <= 0:
except Exception: return
txt_path, part_path, meta_path = self._cache_paths(self._resolved_url) await self.update_progress(min(downloaded, total), commit=True)
txt_path.unlink(missing_ok=True)
part_path.unlink(missing_ok=True) body_path = await self._downloader.download_file(
meta_path.unlink(missing_ok=True) client,
body = await self._download_body_with_resume(client, self._resolved_url) self._resolved_url,
extension=".txt",
progress_callback=on_progress,
)
body = body_path.read_text(encoding="utf-8", errors="replace")
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []
for raw_line in body.splitlines(): for raw_line in body.splitlines():

View File

@@ -14,6 +14,7 @@ from typing import Any
import httpx import httpx
from app.services.collectors.base import BaseCollector from app.services.collectors.base import BaseCollector
from app.services.collectors.downloads import ResumableFileDownloader
class OpenGeoFeedPrefixGeoCollector(BaseCollector): class OpenGeoFeedPrefixGeoCollector(BaseCollector):
@@ -23,21 +24,37 @@ class OpenGeoFeedPrefixGeoCollector(BaseCollector):
frequency_hours = 24 frequency_hours = 24
data_type = "prefix_geography" data_type = "prefix_geography"
fail_on_empty = True fail_on_empty = True
_downloader = ResumableFileDownloader(
cache_namespace="opengeofeed",
default_accept="text/csv,*/*",
)
async def fetch(self) -> list[dict[str, Any]]: async def fetch(self) -> list[dict[str, Any]]:
if not self._resolved_url: if not self._resolved_url:
raise RuntimeError("OpenGeoFeed URL is not configured") raise RuntimeError("OpenGeoFeed URL is not configured")
async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client: async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client:
response = await client.get( remote = await self._downloader.fetch_remote_info(client, self._resolved_url)
total_expected = remote.get("content_length") or 0
if total_expected > 0 and self._current_task and self._db_session:
self._current_task.total_records = total_expected
self._current_task.records_processed = 0
self._current_task.progress = 0.0
await self._db_session.commit()
await self._publish_task_update(force=True)
async def on_progress(downloaded: int, total: int | None) -> None:
if not total or total <= 0:
return
await self.update_progress(min(downloaded, total), commit=True)
body_path = await self._downloader.download_file(
client,
self._resolved_url, self._resolved_url,
headers={ extension=".csv",
"User-Agent": "Planet-Intelligence-System/1.0 (Python/collector)", progress_callback=on_progress,
"Accept": "text/csv,*/*",
},
) )
response.raise_for_status() body = body_path.read_text(encoding="utf-8", errors="replace")
body = response.text
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []
reader = csv.reader(body.splitlines()) reader = csv.reader(body.splitlines())