diff --git a/.gitignore b/.gitignore index 183c527..b7a963a 100644 --- a/.gitignore +++ b/.gitignore @@ -110,6 +110,3 @@ test_data/ TODO.md newsletter.py notion/ - -# Tests -tests/downloaders/ \ No newline at end of file diff --git a/bluemath_tk/downloaders/__init__.py b/bluemath_tk/downloaders/__init__.py index 8580035..d8519c8 100644 --- a/bluemath_tk/downloaders/__init__.py +++ b/bluemath_tk/downloaders/__init__.py @@ -4,6 +4,20 @@ Author: GeoOcean Research Group, Universidad de Cantabria Repository: https://github.com/GeoOcean/BlueMath_tk.git Status: Under development (Working) + +Downloaders for climate and ocean data providers. + +Each provider has its own subpackage with one downloader class and one folder +per product holding the product configuration:: + + downloaders//_downloader.py + downloaders///_config.json + +Providers are imported from their subpackage, so that their optional +dependencies are only needed when used, e.g.:: + + from bluemath_tk.downloaders.csiro import CSIRODownloader + from bluemath_tk.downloaders.copernicus import CopernicusDownloader """ from ._base_downloaders import BaseDownloader diff --git a/bluemath_tk/downloaders/_base_downloaders.py b/bluemath_tk/downloaders/_base_downloaders.py index c2840ea..1fef60b 100644 --- a/bluemath_tk/downloaders/_base_downloaders.py +++ b/bluemath_tk/downloaders/_base_downloaders.py @@ -1,9 +1,37 @@ -from abc import abstractmethod +import json +import os +import re from datetime import datetime -from typing import List, Optional +from typing import Callable, Dict, List, Optional from ..core.models import BlueMathModel from ._download_result import DownloadResult +from ._transfer import PART_SUFFIX, RemoteFileNotFound, retry_call, validate_file + + +def load_config(*path_parts: str) -> dict: + """ + Load a JSON configuration file. + + Parameters + ---------- + *path_parts : str + Parts of the path, joined with os.path.join. + + Returns + ------- + dict + The parsed configuration. + """ + + with open(os.path.join(*path_parts)) as f: + return json.load(f) + + +def _to_snake_case(name: str) -> str: + """Convert a product name to snake case: "OpenData" -> "open_data", "ERA5" -> "era5".""" + + return re.sub(r"(?<=[a-z0-9])(?=[A-Z][a-z])", "_", name).lower() class BaseDownloader(BlueMathModel): @@ -11,9 +39,15 @@ class BaseDownloader(BlueMathModel): Abstract base class for BlueMath downloaders. All downloaders should: - 1. Have a `download_data` method that routes to product-specific methods - 2. Have product-specific methods like `download_data_` - 3. Use DownloadResult to track download status + + 1. Declare their products in the class attribute ``products_configs``, + loaded from ``//_config.json``. + 2. Implement product-specific methods named ``download_data_`` + (snake case), which ``download_data`` routes to, or override + ``download_data`` when all products share the same logic. + 3. Write files through ``fetch_file``, so that every file is downloaded to a + ``.part`` file, validated and then atomically renamed, with retries. + 4. Use DownloadResult to track download status. Attributes ---------- @@ -25,13 +59,21 @@ class BaseDownloader(BlueMathModel): Base path where downloaded files are stored. debug : bool If True, logger is set to DEBUG level. + retries : int + Number of attempts per file. + retry_backoff : float + Base waiting time in seconds between attempts (doubles each attempt). """ + products_configs: Dict[str, dict] = {} + def __init__( self, product: str, base_path_to_download: str, debug: bool = True, + retries: int = 3, + retry_backoff: float = 30.0, ) -> None: """ Initialize the BaseDownloader. @@ -44,11 +86,25 @@ def __init__( The base path to download the data. debug : bool, optional If True, the logger will be set to DEBUG level. Default is True. + retries : int, optional + Number of attempts per file. Default is 3. + retry_backoff : float, optional + Base waiting time in seconds between attempts. Default is 30. + + Raises + ------ + ValueError + If the product is not in ``products_configs`` or an argument has a + wrong type. """ super().__init__() if not isinstance(product, str): raise ValueError("product must be a string") + if product not in self.products_configs: + raise ValueError( + f"Product '{product}' not found. Available: {self.list_products()}" + ) self._product: str = product if not isinstance(base_path_to_download, str): raise ValueError("base_path_to_download must be a string") @@ -56,6 +112,10 @@ def __init__( if not isinstance(debug, bool): raise ValueError("debug must be a boolean") self._debug: bool = debug + if not isinstance(retries, int) or retries < 1: + raise ValueError("retries must be a positive integer") + self.retries: int = retries + self.retry_backoff: float = float(retry_backoff) @property def product(self) -> str: @@ -70,9 +130,29 @@ def debug(self) -> bool: return self._debug @property - @abstractmethod def product_config(self) -> dict: - pass + """ + Configuration of the product, loaded from its JSON file. + + Returns + ------- + dict + Product configuration dictionary. + """ + return self.products_configs[self.product] + + @classmethod + def list_products(cls) -> List[str]: + """ + List the products this downloader supports. + + Returns + ------- + List[str] + List of product names. + """ + + return list(cls.products_configs.keys()) def list_datasets(self) -> List[str]: """ @@ -143,30 +223,131 @@ def finalize_download_result( parts.append(f"{len(result.skipped_files)} skipped") if result.error_files: parts.append(f"{len(result.error_files)} errors") - result.message = f"Download complete: {', '.join(parts)}" + result.message = f"Download complete: {', '.join(parts) or 'nothing to do'}" else: result.message = message return result - @abstractmethod - def download_data(self, *args, **kwargs) -> DownloadResult: + def fetch_file( + self, + output_file: str, + fetch: Callable[[str], None], + result: DownloadResult, + force: bool = False, + dry_run: bool = True, + resumable: bool = False, + validate: bool = True, + ) -> None: + """ + Download a single file safely and record the outcome in ``result``. + + ``fetch`` is called with a temporary ``.part`` path and must + write the whole file there. The ``.part`` file is then validated and + renamed to ``output_file``, so ``output_file`` only ever exists complete. + Failed attempts are retried ``self.retries`` times. + + Parameters + ---------- + output_file : str + Final path of the file. + fetch : Callable[[str], None] + Function that writes the file to the path it receives. + result : DownloadResult + Result to update. + force : bool, optional + Download even if ``output_file`` exists. Default is False. + dry_run : bool, optional + Only record what would be downloaded. Default is True. + resumable : bool, optional + If True, an existing ``.part`` file is kept between attempts so that + ``fetch`` can resume it. Default is False. + validate : bool, optional + Check the downloaded file with validate_file. Default is True. + """ + + if not force and os.path.exists(output_file): + result.add_skipped(output_file, "Already downloaded") + return + + if dry_run: + result.add_skipped(output_file, "Would download (dry run)") + return + + part_file = output_file + PART_SUFFIX + os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True) + + def attempt() -> None: + if not resumable and os.path.exists(part_file): + os.remove(part_file) + fetch(part_file) + if validate: + try: + validate_file(part_file) + except Exception: + os.remove(part_file) + raise + os.replace(part_file, output_file) + + try: + self.logger.debug(f"Downloading: {output_file}") + retry_call( + attempt, + retries=self.retries, + backoff=self.retry_backoff, + logger=self.logger, + ) + result.add_downloaded(output_file) + self.logger.info(f"Downloaded: {output_file}") + + except RemoteFileNotFound as e: + self._remove_quietly(part_file) + result.add_skipped(output_file, f"Not available on server: {e}") + self.logger.info(f"Not available on server: {e}") + + except Exception as e: + if not resumable: + self._remove_quietly(part_file) + result.add_error(output_file, e) + self.logger.error(f"Error downloading {output_file}: {e}") + + @staticmethod + def _remove_quietly(path: str) -> None: + try: + os.remove(path) + except FileNotFoundError: + pass + + def download_data(self, *args, dry_run: bool = True, **kwargs) -> DownloadResult: """ Download data for the product. - Routes to product-specific methods like download_data_(). + Routes to the product-specific method ``download_data_``, where + the product name is converted to snake case (e.g. "OpenData" -> "open_data"). Parameters ---------- *args Arguments passed to product-specific download method. + dry_run : bool, optional + If True, only check what would be downloaded. Default is True. **kwargs - Keyword arguments (e.g., force, dry_run). + Keyword arguments passed to product-specific download method (e.g., force). Returns ------- DownloadResult Result with information about downloaded, skipped, and error files. + + Raises + ------ + ValueError + If there is no download method for the product. """ - pass + method_name = f"download_data_{_to_snake_case(self.product)}" + method = getattr(self, method_name, None) + if method is None: + raise ValueError(f"Download for product {self.product} not supported") + + return method(*args, dry_run=dry_run, **kwargs) diff --git a/bluemath_tk/downloaders/_download_result.py b/bluemath_tk/downloaders/_download_result.py index 94e09e1..0809486 100644 --- a/bluemath_tk/downloaders/_download_result.py +++ b/bluemath_tk/downloaders/_download_result.py @@ -1,3 +1,6 @@ +import json +import os +from collections import Counter from dataclasses import dataclass, field from datetime import datetime from typing import Any, Dict, List, Optional @@ -133,6 +136,96 @@ def add_skipped(self, file_path: str, reason: str = ""): if reason: self.metadata.setdefault("skip_reasons", {})[file_path] = reason + def merge(self, other: "DownloadResult") -> "DownloadResult": + """ + Add the files, errors and metadata of another result to this one. + + Dictionary metadata (e.g. skip reasons) are merged, other values replaced. + + Parameters + ---------- + other : DownloadResult + The result to merge into this one. + + Returns + ------- + DownloadResult + This result, updated. + """ + for file_path in other.downloaded_files: + self.add_downloaded(file_path) + for file_path in other.skipped_files: + if file_path not in self.skipped_files: + self.skipped_files.append(file_path) + for file_path in other.error_files: + if file_path not in self.error_files: + self.error_files.append(file_path) + self.errors.extend(other.errors) + for key, value in other.metadata.items(): + if isinstance(value, dict): + self.metadata.setdefault(key, {}).update(value) + else: + self.metadata[key] = value + return self + + def save(self, file_path: str) -> None: + """ + Save the result as JSON (e.g. to keep a log of automated runs). + + Parameters + ---------- + file_path : str + Path of the JSON file to write. + """ + os.makedirs(os.path.dirname(file_path) or ".", exist_ok=True) + with open(file_path, "w") as f: + json.dump(self.to_dict(), f, indent=2, default=str) + + def summary(self, max_errors: int = 20) -> str: + """ + Multi-line report: counts, skip reasons, size to download and errors. + + Parameters + ---------- + max_errors : int, optional + Maximum number of errors listed. Default is 20. + + Returns + ------- + str + The report. + """ + lines = [ + str(self), + f" downloaded: {len(self.downloaded_files)}", + f" skipped: {len(self.skipped_files)}", + f" errors: {len(self.error_files)}", + ] + reasons = Counter(self.metadata.get("skip_reasons", {}).values()) + for reason, count in reasons.most_common(): + lines.append(f" {count:6d} x {reason}") + + remote_sizes = self.metadata.get("remote_sizes", {}) + pending = [ + path + for path, reason in self.metadata.get("skip_reasons", {}).items() + if "dry run" in reason + ] + if pending and remote_sizes: + total = sum(remote_sizes.get(path, 0) for path in pending) + lines.append( + f" to download: {total / 1e9:,.1f} GB in {len(pending)} files" + ) + + for error in self.errors[:max_errors]: + lines.append(f" ERROR {error['file']}: {error['error']}") + if len(self.errors) > max_errors: + lines.append(f" ... and {len(self.errors) - max_errors} more errors") + if self.duration_seconds is not None: + lines.append(f" duration: {self.duration_seconds / 3600:.2f} h") + + return "\n".join(lines) + def to_dict(self) -> Dict[str, Any]: """Convert the result to a dictionary.""" return { diff --git a/bluemath_tk/downloaders/_transfer.py b/bluemath_tk/downloaders/_transfer.py new file mode 100644 index 0000000..37c4bb9 --- /dev/null +++ b/bluemath_tk/downloaders/_transfer.py @@ -0,0 +1,231 @@ +""" +Low-level transfer helpers shared by all downloaders. + +These functions are deliberately independent from any data provider, so every +downloader gets the same behaviour for retries, resumable HTTP transfers and +file validation. +""" + +import logging +import os +import time +from typing import Callable, Optional, Tuple, Type, TypeVar + +import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + +PART_SUFFIX = ".part" +CHUNK_SIZE = 8 * 1024 * 1024 + +T = TypeVar("T") + + +class RemoteFileNotFound(Exception): + """Raised when a remote file does not exist (e.g. not published yet).""" + + +def retry_call( + func: Callable[[], T], + retries: int = 3, + backoff: float = 30.0, + logger: Optional[logging.Logger] = None, + no_retry: Tuple[Type[BaseException], ...] = (RemoteFileNotFound,), +) -> T: + """ + Call ``func`` until it succeeds, waiting ``backoff * 2**n`` seconds between attempts. + + Parameters + ---------- + func : Callable[[], T] + Function without arguments to call. + retries : int, optional + Maximum number of attempts. Default is 3. + backoff : float, optional + Base waiting time in seconds. Default is 30. + logger : logging.Logger, optional + Logger used to report failed attempts. + no_retry : Tuple[Type[BaseException], ...], optional + Exceptions that are raised immediately, without retrying. + + Returns + ------- + T + Whatever ``func`` returns. + """ + + for attempt in range(1, retries + 1): + try: + return func() + except no_retry: + raise + except Exception as e: + if attempt == retries: + raise + wait = backoff * 2 ** (attempt - 1) + if logger is not None: + logger.warning( + f"Attempt {attempt}/{retries} failed ({type(e).__name__}: {e}). " + f"Retrying in {wait:.0f}s" + ) + time.sleep(wait) + + +def make_session(retries: int = 5, backoff: float = 2.0) -> requests.Session: + """ + Create a requests session that retries connection errors and 429/5xx responses. + + Parameters + ---------- + retries : int, optional + Number of retries at the HTTP level. Default is 5. + backoff : float, optional + urllib3 backoff factor. Default is 2. + + Returns + ------- + requests.Session + Configured session. + """ + + retry = Retry( + total=retries, + backoff_factor=backoff, + status_forcelist=(429, 500, 502, 503, 504), + allowed_methods=("HEAD", "GET"), + raise_on_status=False, + ) + session = requests.Session() + adapter = HTTPAdapter(max_retries=retry) + session.mount("http://", adapter) + session.mount("https://", adapter) + + return session + + +def http_get(session: requests.Session, url: str, **kwargs) -> requests.Response: + """ + GET ``url`` raising RemoteFileNotFound on 404 and HTTPError on other failures. + """ + + kwargs.setdefault("timeout", (30, 300)) + response = session.get(url, **kwargs) + if response.status_code == 404: + response.close() + raise RemoteFileNotFound(url) + response.raise_for_status() + + return response + + +def remote_size(session: requests.Session, url: str) -> Optional[int]: + """ + Size in bytes of a remote file (HEAD request), or None if the server does not say. + + Raises + ------ + RemoteFileNotFound + If the server answers 404, or reports an empty file (THREDDS answers + HEAD requests for missing files with 200 and Content-Length: 0). + """ + + response = session.head(url, allow_redirects=True, timeout=(30, 120)) + if response.status_code == 404: + raise RemoteFileNotFound(url) + response.raise_for_status() + content_length = response.headers.get("Content-Length") + if content_length is None: + return None + if int(content_length) == 0: + raise RemoteFileNotFound(url) + + return int(content_length) + + +def http_download( + session: requests.Session, + url: str, + target: str, + expected_size: Optional[int] = None, + resume: bool = True, +) -> None: + """ + Stream ``url`` into ``target``, resuming a partial ``target`` with a Range request. + + Parameters + ---------- + session : requests.Session + Session used for the request. + url : str + Remote URL. + target : str + Local file to write (usually a ``.part`` file). + expected_size : int, optional + Expected final size in bytes. If given, the download is checked against it. + resume : bool, optional + Whether to continue an existing partial ``target``. Default is True. + + Raises + ------ + RemoteFileNotFound + If the server answers 404. + IOError + If the final size does not match ``expected_size``. + """ + + offset = os.path.getsize(target) if resume and os.path.exists(target) else 0 + if expected_size is not None: + if offset == expected_size: + return + if offset > expected_size: + offset = 0 + + headers = {"Range": f"bytes={offset}-"} if offset else {} + with http_get(session, url, headers=headers, stream=True) as response: + # The server may ignore the Range header and send the full file + if offset and response.status_code != 206: + offset = 0 + with open(target, "ab" if offset else "wb") as f: + for chunk in response.iter_content(chunk_size=CHUNK_SIZE): + f.write(chunk) + + if expected_size is not None: + size = os.path.getsize(target) + if size != expected_size: + raise IOError( + f"Incomplete download of {url}: {size} of {expected_size} bytes" + ) + + +def validate_file(path: str) -> None: + """ + Check that a downloaded file is readable, based on its magic bytes. + + NetCDF/HDF5 files are opened with netCDF4 and GRIB files must end with the + GRIB end marker. Other formats are only checked to be non-empty. + + Raises + ------ + IOError + If the file is empty or cannot be read. + """ + + if os.path.getsize(path) == 0: + raise IOError(f"Empty file: {path}") + + with open(path, "rb") as f: + magic = f.read(8) + + if magic.startswith(b"CDF") or magic.startswith(b"\x89HDF"): + import netCDF4 + + try: + netCDF4.Dataset(path).close() + except Exception as e: + raise IOError(f"Corrupt NetCDF file {path}: {e}") from e + + elif magic.startswith(b"GRIB"): + with open(path, "rb") as f: + f.seek(-4, os.SEEK_END) + if f.read(4) != b"7777": + raise IOError(f"Truncated GRIB file: {path}") diff --git a/bluemath_tk/downloaders/aviso/__init__.py b/bluemath_tk/downloaders/aviso/__init__.py index e69de29..006dcf4 100644 --- a/bluemath_tk/downloaders/aviso/__init__.py +++ b/bluemath_tk/downloaders/aviso/__init__.py @@ -0,0 +1,11 @@ +""" +Project: BlueMath_tk +Sub-Module: downloaders.aviso +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) +""" + +from .aviso_downloader import AvisoDownloader + +__all__ = ["AvisoDownloader"] diff --git a/bluemath_tk/downloaders/aviso/aviso_downloader.py b/bluemath_tk/downloaders/aviso/aviso_downloader.py index f147841..6c6d911 100644 --- a/bluemath_tk/downloaders/aviso/aviso_downloader.py +++ b/bluemath_tk/downloaders/aviso/aviso_downloader.py @@ -1,38 +1,48 @@ +""" +Downloader for AVISO+ altimetry products, served over FTP. + +Products (configs in ``aviso//_config.json``): + +- SWOT: SWOT L2/L3 KaRIn expert products, one directory per cycle. +""" + import ftplib -import json import os +import posixpath from typing import List, Optional -from .._base_downloaders import BaseDownloader +from .._base_downloaders import BaseDownloader, load_config from .._download_result import DownloadResult +_HERE = os.path.dirname(__file__) + class AvisoDownloader(BaseDownloader): """ - Simple downloader for AVISO data. + Downloader for AVISO+ altimetry data, served over FTP (requires an account). - Downloads all available files from the FTP base path specified in the config. + Downloads all .nc files of the requested cycles of a dataset. + The FTP connection is opened on first use, re-opened after failures, and + closed with ``close()`` or when used as a context manager. Examples -------- - >>> downloader = AvisoDownloader( + >>> with AvisoDownloader( ... product="SWOT", ... base_path_to_download="./swot_data", ... username="your_username", ... password="your_password" - ... ) - >>> result = downloader.download_data( - ... dataset="swot-l3-expert", - ... cycles=["cycle_001"], - ... force=False, - ... dry_run=False - ... ) + ... ) as downloader: + ... result = downloader.download_data( + ... dataset="swot-l3-expert", + ... cycles=["cycle_001"], + ... force=False, + ... dry_run=False + ... ) """ products_configs = { - "SWOT": json.load( - open(os.path.join(os.path.dirname(__file__), "SWOT", "SWOT_config.json")) - ) + "SWOT": load_config(_HERE, "SWOT", "SWOT_config.json"), } def __init__( @@ -42,6 +52,7 @@ def __init__( username: str, password: str, debug: bool = True, + **kwargs, ) -> None: """ Initialize the AvisoDownloader. @@ -58,6 +69,8 @@ def __init__( AVISO FTP password. debug : bool, optional If True, sets logger to DEBUG level. Default is True. + **kwargs + Passed to BaseDownloader (e.g., retries, retry_backoff). Raises ------ @@ -66,88 +79,57 @@ def __init__( """ super().__init__( - product=product, base_path_to_download=base_path_to_download, debug=debug + product=product, + base_path_to_download=base_path_to_download, + debug=debug, + **kwargs, ) - self._product_config = self.products_configs.get(product) - if self._product_config is None: - raise ValueError( - f"Product '{product}' not found. Available: {list(self.products_configs.keys())}" - ) - self.set_logger_name( f"AvisoDownloader-{product}", level="DEBUG" if debug else "INFO" ) - # Initialize FTP client - ftp_server = self._product_config.get("ftp_server") - if ftp_server is None: + self._ftp_server = self.product_config.get("ftp_server") + if self._ftp_server is None: raise ValueError("FTP server not found in product configuration") - self._client = ftplib.FTP(ftp_server) - self._client.login(username, password) + self._username = username + self._password = password + self._client: Optional[ftplib.FTP] = None + self._exclude_attributes.extend(["_client", "_password"]) self.logger.info(f"---- AVISO DOWNLOADER INITIALIZED ({product}) ----") - @property - def product_config(self) -> dict: - """ - Product configuration dictionary loaded from config file. + def __enter__(self) -> "AvisoDownloader": + return self - Returns - ------- - dict - Product configuration dictionary. - """ - return self._product_config + def __exit__(self, *exc) -> None: + self.close() @property def client(self) -> ftplib.FTP: """ - FTP client connection (initialized and logged in). + FTP client connection, opened and logged in on first use. Returns ------- ftplib.FTP FTP client instance. """ + if self._client is None: + self._client = ftplib.FTP(self._ftp_server, timeout=120) + self._client.login(self._username, self._password) return self._client - def download_data( - self, - dry_run: bool = True, - *args, - **kwargs, - ) -> DownloadResult: + def close(self) -> None: """ - Download data for the product. - - Routes to product-specific download methods based on the product type. - - Parameters - ---------- - dry_run : bool, optional - If True, only check what would be downloaded without actually downloading. - Default is True. - *args - Arguments passed to product-specific download method. - **kwargs - Keyword arguments passed to product-specific download method. - - Returns - ------- - DownloadResult - Result with information about downloaded, skipped, and error files. - - Raises - ------ - ValueError - If the product is not supported. + Close the FTP connection, if open. """ - - if self.product == "SWOT": - return self.download_data_swot(dry_run=dry_run, *args, **kwargs) - else: - raise ValueError(f"Download for product {self.product} not supported") + if self._client is not None: + try: + self._client.quit() + except Exception: + self._client.close() + self._client = None def download_data_swot( self, @@ -193,144 +175,73 @@ def download_data_swot( dataset_config = self.product_config["datasets"][dataset] ftp_base_path = dataset_config["ftp_base_path"] - result = self.create_download_result() - - try: - if cycles is None: - cycles = dataset_config.get("cycles", []) - if not cycles: - raise ValueError( - f"No cycles specified for dataset '{dataset}' and cycles parameter not provided" - ) - - self.logger.info(f"Downloading dataset: {dataset}, cycles: {cycles}") - - all_downloaded_files = [] - - for cycle in cycles: - files = self._list_all_files_in_cycle(ftp_base_path, cycle) - if not files: - self.logger.warning(f"No files found in cycle {cycle}") - continue - - downloaded_files = self._download_files( - files=files, - dataset=dataset, - ftp_base_path=ftp_base_path, - cycle=cycle, - force=force, - dry_run=dry_run, - result=result, + if cycles is None: + cycles = dataset_config.get("cycles", []) + if not cycles: + raise ValueError( + f"No cycles specified for dataset '{dataset}' and cycles parameter not provided" ) - all_downloaded_files.extend(downloaded_files) - result.downloaded_files = all_downloaded_files - return self.finalize_download_result(result) + self.logger.info(f"Downloading dataset: {dataset}, cycles: {cycles}") + result = self.create_download_result() - except Exception as e: - result.add_error("download_operation", e) - return self.finalize_download_result(result) + for cycle in cycles: + remote_dir = posixpath.join(ftp_base_path, cycle) + try: + files = self._list_nc_files(remote_dir) + except Exception as e: + self.logger.error(f"Error listing {remote_dir}: {e}") + result.add_error(remote_dir, e) + continue + if not files: + self.logger.warning(f"No files found in cycle {cycle}") + continue - def _list_all_files_in_cycle(self, ftp_base_path: str, cycle: str) -> List[str]: - """ - List all .nc files from a cycle directory on FTP server. + for filename in files: + local_path = os.path.join( + self.base_path_to_download, dataset, cycle, filename + ) + remote_path = posixpath.join(remote_dir, filename) - This method navigates to the specified FTP base path and then into the - cycle directory, lists its contents, and filters for files ending with '.nc'. - It assumes the current FTP connection is already logged in. + def fetch(part_file: str, remote_path: str = remote_path) -> None: + try: + with open(part_file, "wb") as f: + self.client.retrbinary(f"RETR {remote_path}", f.write) + except Exception: + # The connection may be broken: reconnect on next attempt + self.close() + raise - Parameters - ---------- - ftp_base_path : str - FTP base path for the dataset (e.g., "/swot_products/l3_karin_nadir/l3_lr_ssh/v1_0/Expert/"). - cycle : str - Cycle directory name (e.g., "cycle_001"). + self.fetch_file(local_path, fetch, result, force=force, dry_run=dry_run) - Returns - ------- - List[str] - List of .nc filenames (without path) found in the cycle directory. - """ + return self.finalize_download_result(result) - files = [] - self._client.cwd(ftp_base_path) - self._client.cwd(cycle) - items = [] - self._client.retrlines("LIST", items.append) - for item in items: - parts = item.split() - if len(parts) >= 9: - name = " ".join(parts[8:]) - if name.endswith(".nc"): - files.append(name) - - return files - - def _download_files( - self, - files: List[str], - dataset: str, - ftp_base_path: str, - cycle: str, - force: bool, - dry_run: bool, - result: DownloadResult, - ) -> List[str]: + def _list_nc_files(self, remote_dir: str) -> List[str]: """ - Download all files from the list. - - Files are saved to: base_path_to_download/dataset/cycle/filename.nc + List the .nc files (names only) in a directory of the FTP server. Parameters ---------- - files : List[str] - List of filenames to download (without path). - dataset : str - Dataset name (used in local path, e.g., "swot-l3-expert"). - ftp_base_path : str - FTP base path for the dataset (e.g., "/swot_products/l3_karin_nadir/l3_lr_ssh/v1_0/Expert/"). - cycle : str - Cycle directory name (used in local path, e.g., "cycle_001"). - force : bool - Force re-download even if file exists. - dry_run : bool - If True, only simulate download. - result : DownloadResult - Download result object to update. + remote_dir : str + Absolute FTP directory (e.g., "/swot_products/.../Expert/cycle_001"). Returns ------- List[str] - List of local file paths for successfully downloaded files only. + Sorted list of .nc filenames. """ - downloaded_files = [] - - for filename in files: - local_path = os.path.join( - self.base_path_to_download, dataset, cycle, filename - ) - - if not force and os.path.exists(local_path): - result.add_skipped(local_path, "Already downloaded") - continue - - if dry_run: - result.add_skipped(local_path, f"Would download {filename} (dry run)") - continue - - try: - os.makedirs(os.path.dirname(local_path), exist_ok=True) - self._client.cwd(ftp_base_path) - self._client.cwd(cycle) - with open(local_path, "wb") as f: - self._client.retrbinary(f"RETR {filename}", f.write) - result.add_downloaded(local_path) - self.logger.info(f"Downloaded: {filename} -> {local_path}") - downloaded_files.append(local_path) - - except Exception as e: - result.add_error(local_path, e) - self.logger.error(f"Error downloading {filename}: {e}") - - return downloaded_files + try: + names = self.client.nlst(remote_dir) + except ftplib.error_perm as e: + # Some servers answer "550 No files found" for empty directories + if str(e).startswith("550"): + return [] + raise + except (ftplib.error_temp, OSError, EOFError): + self.close() + names = self.client.nlst(remote_dir) + + return sorted( + {posixpath.basename(name) for name in names if name.endswith(".nc")} + ) diff --git a/bluemath_tk/downloaders/copernicus/ERA5/ERA5_download.py b/bluemath_tk/downloaders/copernicus/ERA5/ERA5_download.py deleted file mode 100755 index aba80c3..0000000 --- a/bluemath_tk/downloaders/copernicus/ERA5/ERA5_download.py +++ /dev/null @@ -1,97 +0,0 @@ -#!/usr/bin/env python - -# MIGHT NOT WORK perfectly due to changes in CopernicusDownloader! - -""" -Usage: era5_download.py YEARS MONTHS [--var=] [--base-path=] [--token=] [--debug] [--check] [--force] -era5_download.py (-h | --help | --version) -era5_download.py --list-vars -era5_download.py --list-datasets -era5_download.py --show-info - -This script downloads ERA5 reanalysis data for multiple variables and saves them as NetCDF files. - -Arguments: - YEARS Years of the data to download, separated by commas - To download a single year, use the format YYYY - To download a range of years, use the format YYYY-YYYY - MONTHS Months of the data to download, separated by commas - To download a single month, use the format MM - To download a range of months, use the format MM-MM - -Options: - --var= List of variable values separated by commas. - --base-path= Base path to save the data [default: /lustre/geocean/DATA/Copernicus/]. - --token= Token for the Climate Data Store API. - --debug Debug mode (it is recommended to use this mode). - --check If just file checking is required. - --force Overwrite the data in case it exists. - --list-vars Return a list with the available variables. - --list-datasets Return a list with the available datasets. - --show-info Show information about the variables. - -Examples: - era5_download.py 2014,2016 1,10 --var tp,sst --debug - era5_download.py 2014,2016 1,10 --var tp,sst --debug --check - era5_download.py --list-vars - era5_download.py --list-datasets - era5_download.py --show-info -""" - -from docopt import docopt -from bluemath_tk.downloaders.copernicus.copernicus_downloader import ( - CopernicusDownloader, -) - -# Parse the command line arguments -args = docopt(__doc__) -print(args) - -# Create the CopernicusDownloader object -copernicus_downloader = CopernicusDownloader( - product="ERA5", - token=args["--token"], - base_path_to_download=args["--base-path"] or "/lustre/geocean/DATA/Copernicus/", - debug=args["--debug"], - check=args["--check"], -) - -if args["--list-vars"]: - print(copernicus_downloader.list_variables()) - exit() -if args["--list-datasets"]: - print(copernicus_downloader.list_datasets()) - exit() -if args["--show-info"]: - print(copernicus_downloader.show_markdown_table()) - exit() - -years = ( - args["YEARS"].split(",") - if "-" not in args["YEARS"] - else list( - range(int(args["YEARS"].split("-")[0]), int(args["YEARS"].split("-")[1]) + 1) - ) -) -months = ( - args["MONTHS"].split(",") - if "-" not in args["MONTHS"] - else list( - range(int(args["MONTHS"].split("-")[0]), int(args["MONTHS"].split("-")[1]) + 1) - ) -) -# Ocean variables are used by default, as right now, DestinE program is used for atmosphere variables -variables = ( - args["--var"].split(",") - if args["--var"] - else copernicus_downloader.list_variables(type="ocean") -) - -output = copernicus_downloader.download_data( - variables=variables, - years=years, - months=months, - force=args["--force"], -) - -print(output) diff --git a/bluemath_tk/downloaders/copernicus/ERA5/README.md b/bluemath_tk/downloaders/copernicus/ERA5/README.md deleted file mode 100644 index d7ca0a1..0000000 --- a/bluemath_tk/downloaders/copernicus/ERA5/README.md +++ /dev/null @@ -1,149 +0,0 @@ -# Reanalysis ERA5 - -For updated documentation please go [here](https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation). - -ERA5 is the fifth generation ECMWF atmospheric reanalysis of the global climate covering the period from January 1940 to present. ERA5 is produced by the Copernicus Climate Change Service (C3S) at ECMWF. - -ERA5 provides hourly estimates for a large number of atmospheric, ocean-wave and land-surface quantities. An uncertainty estimate is sampled by an underlying 10-member ensemble at three-hourly intervals. - -**There is a script scheduled to run periodically to download the latest month's data.** - -## Dataset Characteristics - -The Reanalysis ERA5 single levels dataset has the following characteristics: - -- Data type: Gridded -- Projection: Regular latitude-longitude grid -- Horizontal coverage: Global -- Horizontal resolution: - - Reanalysis: 0.25° x 0.25° (atmosphere), 0.5° x 0.5° (ocean waves) - - Mean, spread and members: 0.5° x 0.5° (atmosphere), 1° x 1° (ocean waves) -- Temporal coverage: 1940 to present -- Temporal resolution: Hourly -- File format: NetCDF -- Update frequency: Daily - -* geopotencial at 500Hpa is obtained from Reanalysis ERA5 pressure levels: https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-pressure-levels?tab=overview - -## Downloaded Dataset Variables - -These are the variables currently available for download. If a new variable is required, it MUST be added to the `ERA5_config.json` file, and then to the code or crontab downloading the files. - -| name | long_name | units | type | -| ------- | --------------------------------------------------- | ------------- | ----- | -| swh | Significant height of combined wind waves and swell | m | ocean | -| pp1d | Peak wave period | s | ocean | -| mwp | Mean wave period | s | ocean | -| mwd | Mean wave direction | degrees | ocean | -| shww | Significant height of wind waves | m | ocean | -| mpww | Mean period of wind waves | s | ocean | -| mdww | Mean direction of wind waves | degrees | ocean | -| dwww | Wave spectral directional width for wind waves | degrees | ocean | -| p140121 | Significant wave height of first swell partition | m | ocean | -| p140124 | Significant wave height of second swell partition | m | ocean | -| p140127 | Significant wave height of third swell partition | m | ocean | -| p140123 | Mean wave period of first swell partition | s | ocean | -| p140126 | Mean wave period of second swell partition | s | ocean | -| p140129 | Mean wave period of third swell partition | s | ocean | -| p140122 | Mean wave direction of first swell partition | degrees | ocean | -| p140125 | Mean wave direction of second swell partition | degrees | ocean | -| p140128 | Mean wave direction of third swell partition | degrees | ocean | -| mzl | Mean zero-crossing wave period | s | ocean | -| mbathy | Model bathymetry | m | ocean | -| wsk | Wave spectral kurtosis | dimensionless | ocean | -| wsp | Wave spectral peakedness | dimensionless | ocean | -| wss | Wave spectral skewness | dimensionless | ocean | -| wsdw | Wave spectral directional width | degrees | ocean | - -These variables represent different atmospheric, ocean-wave, and land-surface quantities that are provided by the ERA5 dataset. Each variable has a corresponding path that specifies where the data is located within the dataset. - -For more information and to access the dataset, click [here](https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-single-levels?tab=overview). - -## Scripts used to download data - -To download the Reanalysis ERA5 single levels dataset, you can use the following scripts: - -1. Main Python script: `copernicus_downloader.py`. - -```python -from bluemath_tk.downloaders.copernicus.copernicus_downloader import CopernicusDownloader - -copernicus_downloader = CopernicusDownloader( - product="ERA5", - base_path_to_download="/path/to/Copernicus/", # Will be created if not available - token=None, - check=True, -) -result = copernicus_downloader.download_data_era5( - variables=["swh"], - years=["2020"], - months=["01", "03"], -) -print(result) -``` - -2. Python script: `ERA5_download.py`. - -This script downloads ERA5 reanalysis data for multiple variables and saves them as NetCDF files. - -``` -Usage: era5_download.py YEARS MONTHS [--var=] [--base-path=] [--token=] [--debug] [--check] [--force] -era5_download.py (-h | --help | --version) -era5_download.py --list-vars -era5_download.py --list-datasets -era5_download.py --show-info - -This script downloads ERA5 reanalysis data for multiple variables and saves them as NetCDF files. - -Arguments: - YEARS Years of the data to download, separated by commas - To download a single year, use the format YYYY - To download a range of years, use the format YYYY-YYYY - MONTHS Months of the data to download, separated by commas - To download a single month, use the format MM - To download a range of months, use the format MM-MM - -Options: - --var= List of variable values separated by commas. - --base-path= Base path to save the data [default: /lustre/geocean/DATA/Copernicus/]. - --token= Token for the Climate Data Store API. - --debug Debug mode (it is recommended to use this mode). - --check If just file checking is required. - --force Overwrite the data in case it exists. - --list-vars Return a list with the available variables. - --list-datasets Return a list with the available datasets. - --show-info Show information about the variables. - -Examples: - era5_download.py 2014,2016 1,10 --var tp,sst --debug - era5_download.py 2014,2016 1,10 --var tp,sst --debug --check - era5_download.py --list-vars - era5_download.py --list-datasets - era5_download.py --show-info -``` - -3. Bash scripts: `bash_ERA5_download.sh and qsub_ERA5_download.sh` in [gitlab](https://gitlab.com/geoocean/datahub/datahub-scripts/-/tree/main/ERA5/launchers?ref_type=heads). - -This script creates and submits jobs to download several datasets simultaneously. - -```Exception: Requests using the API temporally limited to 20 to restrict the activity of abusive users. Please visit copernicus-support.ecmwf.int for further information.``` - -For more information on how to use the scripts and download the Reanalysis ERA5 single levels dataset, please refer to the [documentation](https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-single-levels?tab=overview). - -## Issues with data download limits - -To ensure the protection of CDS resources, limits have been set for data downloads. You can find more information about these limits [here](https://cds.climate.copernicus.eu/live/limits). - -To adhere to these limits, we have developed the `download_ERA5.sh` script and configured the cron of the user responsible for data downloads to initiate periodic downloads. To prevent user-specific limits, we have set up cron to periodically update the credentials of the user performing the data downloads. - -``` -# VA TODO 2 HORAS ANTES. Si quieremos ejecutar a las 10, tenemos que poner las 8. - -0 0 * * * /usr/bin/cp /home/grupos/geocean/valvanuz/.cds/laura /home/grupos/geocean/valvanuz/.cdsapirc -0 6 * * * /usr/bin/cp /home/grupos/geocean/valvanuz/.cds/gmail1 /home/grupos/geocean/valvanuz/.cdsapirc -0 12 * * * /usr/bin/cp /home/grupos/geocean/valvanuz/.cds/valva /home/grupos/geocean/valvanuz/.cdsapirc -0 18 * * * /usr/bin/cp /home/grupos/geocean/valvanuz/.cds/isra /home/grupos/geocean/valvanuz/.cdsapirc -30 10 21 * * /usr/bin/bash -l /home/grupos/geocean/valvanuz/data_lustre/datahub-scripts/ERA5/download_ERA5.sh 1941 1941 -0 12 21 * * /usr/bin/bash -l /home/grupos/geocean/valvanuz/data_lustre/datahub-scripts/ERA5/download_ERA5.sh 1940 1945 -0 18 21 * * /usr/bin/bash -l /home/grupos/geocean/valvanuz/data_lustre/datahub-scripts/ERA5/download_ERA5.sh 1946 1950 -``` \ No newline at end of file diff --git a/bluemath_tk/downloaders/copernicus/__init__.py b/bluemath_tk/downloaders/copernicus/__init__.py index e69de29..4a30475 100644 --- a/bluemath_tk/downloaders/copernicus/__init__.py +++ b/bluemath_tk/downloaders/copernicus/__init__.py @@ -0,0 +1,11 @@ +""" +Project: BlueMath_tk +Sub-Module: downloaders.copernicus +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) +""" + +from .copernicus_downloader import CopernicusDownloader + +__all__ = ["CopernicusDownloader"] diff --git a/bluemath_tk/downloaders/copernicus/copernicus_downloader.py b/bluemath_tk/downloaders/copernicus/copernicus_downloader.py index 9e7c50c..f5ae5d0 100644 --- a/bluemath_tk/downloaders/copernicus/copernicus_downloader.py +++ b/bluemath_tk/downloaders/copernicus/copernicus_downloader.py @@ -1,23 +1,47 @@ -import json +""" +Downloader for the Copernicus Climate Data Store (CDS), through ``cdsapi``. + +Products (configs in ``copernicus//_config.json``): + +- ERA5: global reanalysis (single levels, pressure levels, wave spectra). +- CERRA: regional reanalysis for Europe (single levels). + +The CDS key is read from ~/.cdsapirc or the CDSAPI_KEY environment variable, +unless it is passed explicitly. +""" + +import calendar +import copy import os -from typing import Any, Dict, List, Optional +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, Iterator, List, Optional, Tuple import cdsapi +import requests -from .._base_downloaders import BaseDownloader +from .._base_downloaders import BaseDownloader, load_config from .._download_result import DownloadResult +_HERE = os.path.dirname(__file__) +ALL_MONTHS = [f"{month:02d}" for month in range(1, 13)] + class CopernicusDownloader(BaseDownloader): """ - Simple downloader for Copernicus Climate Data Store. + Downloader for the Copernicus Climate Data Store (CDS). + + One file is written per variable and year when all 12 months are requested + ({nc_name}_{YYYY}.nc), and one file per variable and month otherwise + ({nc_name}_{YYYY}_{MM}.nc), so the same data never ends up in two files. + Use file_period="month" to always write monthly files. Examples -------- >>> downloader = CopernicusDownloader( ... product="ERA5", ... base_path_to_download="./copernicus_data", - ... token="your_token" + ... api_key="your_api_key" ... ) >>> result = downloader.download_data( ... variables=["swh"], @@ -29,20 +53,18 @@ class CopernicusDownloader(BaseDownloader): """ products_configs = { - "ERA5": json.load( - open(os.path.join(os.path.dirname(__file__), "ERA5", "ERA5_config.json")) - ), - "CERRA": json.load( - open(os.path.join(os.path.dirname(__file__), "CERRA", "CERRA_config.json")) - ), + "ERA5": load_config(_HERE, "ERA5", "ERA5_config.json"), + "CERRA": load_config(_HERE, "CERRA", "CERRA_config.json"), } def __init__( self, product: str, base_path_to_download: str, - api_key: str, + api_key: str = None, + max_workers: int = 1, debug: bool = True, + **kwargs, ) -> None: """ Initialize the CopernicusDownloader. @@ -53,10 +75,16 @@ def __init__( The product to download data from (e.g., "ERA5", "CERRA"). base_path_to_download : str Base path where downloaded files will be stored. - api_key : str - Copernicus CDS API key. + api_key : str, optional + Copernicus CDS API key. If None, it is read by cdsapi from the + CDSAPI_KEY environment variable or ~/.cdsapirc. Default is None. + max_workers : int, optional + Number of CDS requests submitted in parallel. CDS queues requests, + so a few parallel requests reduce the waiting time a lot. Default is 1. debug : bool, optional If True, sets logger to DEBUG level. Default is True. + **kwargs + Passed to BaseDownloader (e.g., retries, retry_backoff). Raises ------ @@ -65,50 +93,51 @@ def __init__( """ super().__init__( - product=product, base_path_to_download=base_path_to_download, debug=debug + product=product, + base_path_to_download=base_path_to_download, + debug=debug, + **kwargs, ) - self._product_config = self.products_configs.get(product) - if self._product_config is None: - raise ValueError( - f"Product '{product}' not found. Available: {list(self.products_configs.keys())}" - ) - self.set_logger_name( f"CopernicusDownloader-{product}", level="DEBUG" if debug else "INFO" ) # Initialize CDS client - server_url = self._product_config.get("url") + server_url = self.product_config.get("url") if server_url is None: raise ValueError("Server URL not found in product configuration") - self._client = cdsapi.Client(url=server_url, key=api_key, debug=self.debug) + self._server_url = server_url + self._api_key = api_key + self.max_workers = max_workers + self._local = threading.local() + self._exclude_attributes.extend(["_local", "_api_key"]) + # Create a client now so that missing credentials fail early + self.client self.logger.info(f"---- COPERNICUS DOWNLOADER INITIALIZED ({product}) ----") - @property - def product_config(self) -> dict: - """ - Product configuration dictionary loaded from config file. - - Returns - ------- - dict - Product configuration dictionary. - """ - return self._product_config - @property def client(self) -> cdsapi.Client: """ - CDS API client (initialized with API key). + CDS API client of the current thread (initialized with API key). Returns ------- cdsapi.Client CDS API client instance. """ - return self._client + if not hasattr(self._local, "client"): + # Own session per thread (cdsapi shares one by default), and no + # progress bars, which flood the logs of batch jobs + self._local.client = cdsapi.Client( + url=self._server_url, + key=self._api_key, + debug=self.debug, + progress=False, + session=requests.Session(), + ) + return self._local.client def list_variables(self, type: str = None) -> List[str]: """ @@ -125,54 +154,15 @@ def list_variables(self, type: str = None) -> List[str]: List of variable names. """ - if type == "ocean": + if type is not None: return [ var_name for var_name, var_info in self.product_config["variables"].items() - if var_info["type"] == "ocean" + if var_info["type"] == type ] return list(self.product_config["variables"].keys()) - def download_data( - self, - dry_run: bool = True, - *args, - **kwargs, - ) -> DownloadResult: - """ - Download data for the product. - - Routes to product-specific download methods based on the product type. - - Parameters - ---------- - dry_run : bool, optional - If True, only check what would be downloaded without actually downloading. - Default is True. - *args - Arguments passed to product-specific download method. - **kwargs - Keyword arguments passed to product-specific download method. - - Returns - ------- - DownloadResult - Result with information about downloaded, skipped, and error files. - - Raises - ------ - ValueError - If the product is not supported. - """ - - if self.product == "ERA5": - return self.download_data_era5(dry_run=dry_run, *args, **kwargs) - elif self.product == "CERRA": - return self.download_data_cerra(dry_run=dry_run, *args, **kwargs) - else: - raise ValueError(f"Download for product {self.product} not supported") - def download_data_era5( self, variables: List[str], @@ -184,6 +174,7 @@ def download_data_era5( product_type: str = "reanalysis", data_format: str = "netcdf", download_format: str = "unarchived", + file_period: str = "auto", force: bool = False, dry_run: bool = True, ) -> DownloadResult: @@ -216,6 +207,10 @@ def download_data_era5( Data format. Default is "netcdf". download_format : str, optional Download format. Default is "unarchived". + file_period : str, optional + "month" for one file per month, "year" for one file per year, or + "auto" (yearly files if all 12 months are requested, monthly otherwise). + Default is "auto". force : bool, optional Force re-download even if file exists. Default is False. dry_run : bool, optional @@ -232,58 +227,55 @@ def download_data_era5( If years or months are empty lists. """ - if not isinstance(variables, list) or len(variables) == 0: - variables = list(self.product_config["variables"].keys()) - if not isinstance(years, list) or len(years) == 0: - raise ValueError("Years must be a non-empty list") - years = [f"{int(year):04d}" for year in years] - if not isinstance(months, list) or len(months) == 0: - raise ValueError("Months must be a non-empty list") - months = [f"{int(month):02d}" for month in months] - last_month = months[-1] + variables, years, months = self._check_inputs(variables, years, months) if days is None: days = [f"{day:02d}" for day in range(1, 32)] if times is None: times = [f"{hour:02d}:00" for hour in range(24)] - result = self.create_download_result() - - # Prepare download tasks - download_tasks = [] + tasks = [] for variable in variables: - for year in years: - task = self._prepare_era5_download_task( - variable=variable, - year=year, - months=months, - days=days, - times=times, - area=area, - product_type=product_type, - data_format=data_format, - download_format=download_format, - last_month=last_month, + variable_config, dataset_config = self._get_variable_configs(variable) + if variable_config is None: + continue + for year, period_months, period in self._split_periods( + years, months, file_period + ): + request = copy.deepcopy(dataset_config["template"]) + if variable_config["dataset"] == "reanalysis-era5-complete": + last_day = calendar.monthrange(int(year), int(period_months[-1]))[1] + request["date"] = ( + f"{year}-{period_months[0]}-01/to/" + f"{year}-{period_months[-1]}-{last_day:02d}" + ) + if area is not None: + request["area"] = "/".join([str(coord) for coord in area]) + else: + request["variable"] = variable_config["cds_name"] + request["year"] = year + request["month"] = period_months + request["day"] = days + request["time"] = times + request["product_type"] = product_type + request["data_format"] = data_format + request["download_format"] = download_format + if area is not None: + request["area"] = area + + task = self._make_task( + variable, + variable_config, + dataset_config, + request, + year, + months, + period, + product_type, ) if task is not None: - download_tasks.append(task) + tasks.append(task) - if not download_tasks: - return self.finalize_download_result( - result, "No valid download tasks found" - ) - - self.logger.info(f"Prepared {len(download_tasks)} download tasks") - - # Download files sequentially - for task in download_tasks: - task_result = self._download_single_file(task, force=force, dry_run=dry_run) - if isinstance(task_result, DownloadResult): - result.downloaded_files.extend(task_result.downloaded_files) - result.skipped_files.extend(task_result.skipped_files) - result.error_files.extend(task_result.error_files) - result.errors.extend(task_result.errors) - - return self.finalize_download_result(result) + return self._run_tasks(tasks, force=force, dry_run=dry_run) def download_data_cerra( self, @@ -297,6 +289,7 @@ def download_data_cerra( data_type: List[str] = None, product_type: str = "analysis", data_format: str = "netcdf", + file_period: str = "auto", force: bool = False, dry_run: bool = True, ) -> DownloadResult: @@ -331,6 +324,10 @@ def download_data_cerra( Product type (e.g., "analysis", "forecast"). Default is "analysis". data_format : str, optional Data format. Default is "netcdf". + file_period : str, optional + "month" for one file per month, "year" for one file per year, or + "auto" (yearly files if all 12 months are requested, monthly otherwise). + Default is "auto". force : bool, optional Force re-download even if file exists. Default is False. dry_run : bool, optional @@ -347,348 +344,222 @@ def download_data_cerra( If years or months are empty lists. """ - if not isinstance(variables, list) or len(variables) == 0: - variables = list(self.product_config["variables"].keys()) - if not isinstance(years, list) or len(years) == 0: - raise ValueError("Years must be a non-empty list") - years = [f"{int(year):04d}" for year in years] - if not isinstance(months, list) or len(months) == 0: - raise ValueError("Months must be a non-empty list") - months = [f"{int(month):02d}" for month in months] - last_month = months[-1] + variables, years, months = self._check_inputs(variables, years, months) if days is None: days = [f"{day:02d}" for day in range(1, 32)] if times is None: - times = [ - "00:00", - "03:00", - "06:00", - "09:00", - "12:00", - "15:00", - "18:00", - "21:00", - ] + times = [f"{hour:02d}:00" for hour in range(0, 24, 3)] if data_type is None: data_type = ["reanalysis"] - result = self.create_download_result() - - # Prepare download tasks - download_tasks = [] + tasks = [] for variable in variables: - for year in years: - task = self._prepare_cerra_download_task( - variable=variable, - year=year, - months=months, - days=days, - times=times, - area=area, - level_type=level_type, - data_type=data_type, - product_type=product_type, - data_format=data_format, - last_month=last_month, + variable_config, dataset_config = self._get_variable_configs(variable) + if variable_config is None: + continue + for year, period_months, period in self._split_periods( + years, months, file_period + ): + request = copy.deepcopy(dataset_config["template"]) + request["variable"] = [variable_config["cds_name"]] + request["level_type"] = level_type + request["data_type"] = data_type + request["product_type"] = product_type + request["year"] = [year] + request["month"] = period_months + request["day"] = days + request["time"] = times + request["data_format"] = data_format + if area is not None: + request["area"] = area + + task = self._make_task( + variable, + variable_config, + dataset_config, + request, + year, + months, + period, + product_type, ) if task is not None: - download_tasks.append(task) + tasks.append(task) - if not download_tasks: - return self.finalize_download_result( - result, "No valid download tasks found" - ) + return self._run_tasks(tasks, force=force, dry_run=dry_run) - self.logger.info(f"Prepared {len(download_tasks)} download tasks") + def _check_inputs( + self, variables: List[str], years: List[str], months: List[str] + ) -> Tuple[List[str], List[str], List[str]]: + """ + Validate and format variables, years ("YYYY") and months ("MM"). - # Download files sequentially - for task in download_tasks: - task_result = self._download_single_file(task, force=force, dry_run=dry_run) - if isinstance(task_result, DownloadResult): - result.downloaded_files.extend(task_result.downloaded_files) - result.skipped_files.extend(task_result.skipped_files) - result.error_files.extend(task_result.error_files) - result.errors.extend(task_result.errors) + Any sequence is accepted (list, tuple, range, array). + """ - return self.finalize_download_result(result) + variables = [] if variables is None else list(variables) + if len(variables) == 0: + variables = list(self.product_config["variables"].keys()) + years = [] if years is None or isinstance(years, str) else list(years) + if len(years) == 0: + raise ValueError("Years must be a non-empty list") + months = [] if months is None or isinstance(months, str) else list(months) + if len(months) == 0: + raise ValueError("Months must be a non-empty list") + years = [f"{int(year):04d}" for year in years] + months = sorted({f"{int(month):02d}" for month in months}) - def _prepare_era5_download_task( - self, - variable: str, - year: str, - months: List[str], - days: List[str], - times: List[str], - area: Optional[List[float]], - product_type: str, - data_format: str, - download_format: str, - last_month: str, - ) -> Optional[Dict[str, Any]]: - """ - Prepare a download task for ERA5. + return variables, years, months - Creates a task dictionary with all necessary information for downloading - a single variable for a single year. + @staticmethod + def _split_periods( + years: List[str], months: List[str], file_period: str = "auto" + ) -> Iterator[Tuple[str, List[str], str]]: + """ + Yield (year, months, period label) for each file to download. - Parameters - ---------- - variable : str - Variable name. - year : str - Year (formatted as "YYYY"). - months : List[str] - List of months (formatted as "MM"). - days : List[str] - List of days (formatted as "DD"). - times : List[str] - List of times (formatted as "HH:MM"). - area : Optional[List[float]] - Geographic area as [north, west, south, east] or None. - product_type : str - Product type. - data_format : str - Data format. - download_format : str - Download format. - last_month : str - Last month in the list (used for date range formatting). + With file_period="auto", a full year gives one file per year, + otherwise one file per month. + """ - Returns - ------- - Optional[Dict[str, Any]] - Task dictionary with download information, or None if configuration is invalid. + if file_period not in ("auto", "month", "year"): + raise ValueError("file_period must be 'auto', 'month' or 'year'") + + for year in years: + if file_period == "year" or ( + file_period == "auto" and months == ALL_MONTHS + ): + yield year, months, year + else: + for month in months: + yield year, [month], f"{year}_{month}" + + def _get_variable_configs( + self, variable: str + ) -> Tuple[Optional[dict], Optional[dict]]: + """ + Get the variable and dataset configurations, or (None, None) if not found. """ variable_config = self.product_config["variables"].get(variable) if variable_config is None: self.logger.error(f"Variable {variable} not found in configuration") - return None + return None, None - variable_dataset = self.product_config["datasets"].get( - variable_config["dataset"] - ) - if variable_dataset is None: + dataset_config = self.product_config["datasets"].get(variable_config["dataset"]) + if dataset_config is None: self.logger.error( f"Dataset {variable_config['dataset']} not found in configuration" ) - return None - - template_for_variable = variable_dataset["template"].copy() - if variable == "spectra": - template_for_variable["date"] = ( - f"{year}-{months[0]}-01/to/{year}-{months[-1]}-31" - ) - if area is not None: - template_for_variable["area"] = "/".join([str(coord) for coord in area]) - else: - template_for_variable["variable"] = variable_config["cds_name"] - template_for_variable["year"] = year - template_for_variable["month"] = months - template_for_variable["day"] = days - template_for_variable["time"] = times - template_for_variable["product_type"] = product_type - template_for_variable["data_format"] = data_format - template_for_variable["download_format"] = download_format - if area is not None: - template_for_variable["area"] = area - - # Check mandatory fields - for mandatory_field in variable_dataset["mandatory_fields"]: - if template_for_variable.get(mandatory_field) is None: - try: - template_for_variable[mandatory_field] = variable_config[ - mandatory_field - ] - except KeyError: - self.logger.error( - f"Mandatory field {mandatory_field} not found for {variable}" - ) - return None + return None, None - # Create output file path - output_nc_file = os.path.join( - self.base_path_to_download, - self.product, - variable_config["dataset"], - variable_config["type"], - product_type, - variable_config["cds_name"], - f"{variable_config['nc_name']}_{year}_{'_'.join(months)}.nc", - ) + return variable_config, dataset_config - return { - "variable": variable, - "year": year, - "variable_config": variable_config, - "variable_dataset": variable_dataset, - "template": template_for_variable, - "output_file": output_nc_file, - "last_month": last_month, - } - - def _prepare_cerra_download_task( + def _make_task( self, variable: str, + variable_config: dict, + dataset_config: dict, + request: dict, year: str, months: List[str], - days: List[str], - times: List[str], - area: Optional[List[float]], - level_type: str, - data_type: List[str], + period: str, product_type: str, - data_format: str, - last_month: str, ) -> Optional[Dict[str, Any]]: """ - Prepare a download task for CERRA. + Check the mandatory fields of a request and build its download task. - Creates a task dictionary with all necessary information for downloading - a single variable for a single year. - - Parameters - ---------- - variable : str - Variable name. - year : str - Year (formatted as "YYYY"). - months : List[str] - List of months (formatted as "MM"). - days : List[str] - List of days (formatted as "DD"). - times : List[str] - List of times (formatted as "HH:MM"). - area : Optional[List[float]] - Geographic area as [north, west, south, east] or None. - level_type : str - Level type. - data_type : List[str] - Data type list. - product_type : str - Product type. - data_format : str - Data format. - last_month : str - Last month in the list (used for date range formatting). + Missing mandatory fields are taken from the variable configuration. Returns ------- Optional[Dict[str, Any]] - Task dictionary with download information, or None if configuration is invalid. + Task with dataset, request, output_file and legacy_file, + or None if a mandatory field is missing. """ - variable_config = self.product_config["variables"].get(variable) - if variable_config is None: - self.logger.error(f"Variable {variable} not found in configuration") - return None - - variable_dataset = self.product_config["datasets"].get( - variable_config["dataset"] - ) - if variable_dataset is None: - self.logger.error( - f"Dataset {variable_config['dataset']} not found in configuration" - ) - return None - - template_for_variable = variable_dataset["template"].copy() - template_for_variable["variable"] = [variable_config["cds_name"]] - template_for_variable["level_type"] = level_type - template_for_variable["data_type"] = data_type - template_for_variable["product_type"] = product_type - template_for_variable["year"] = [year] - template_for_variable["month"] = months - template_for_variable["day"] = days - template_for_variable["time"] = times - template_for_variable["data_format"] = data_format - - if area is not None: - template_for_variable["area"] = area - - # Check mandatory fields - for mandatory_field in variable_dataset["mandatory_fields"]: - if template_for_variable.get(mandatory_field) is None: - self.logger.error( - f"Mandatory field {mandatory_field} not found for {variable}" - ) - return None + for mandatory_field in dataset_config["mandatory_fields"]: + if request.get(mandatory_field) is None: + if mandatory_field not in variable_config: + self.logger.error( + f"Mandatory field {mandatory_field} not found for {variable}" + ) + return None + request[mandatory_field] = variable_config[mandatory_field] - # Create output file path - output_nc_file = os.path.join( + file_format = request.get("data_format", request.get("format", "netcdf")) + extension = "grib" if file_format == "grib" else "nc" + output_dir = os.path.join( self.base_path_to_download, self.product, variable_config["dataset"], variable_config["type"], product_type, variable_config["cds_name"], - f"{variable_config['nc_name']}_{year}_{'_'.join(months)}.nc", ) + nc_name = variable_config["nc_name"] + + # Files downloaded before the one-file-per-period naming + legacy_file = None + if months == ALL_MONTHS: + legacy_file = os.path.join( + output_dir, f"{nc_name}_{year}_{'_'.join(months)}.{extension}" + ) return { "variable": variable, - "year": year, - "variable_config": variable_config, - "template": template_for_variable, - "last_month": last_month, - "output_file": output_nc_file, + "dataset": variable_config["dataset"], + "request": request, + "output_file": os.path.join(output_dir, f"{nc_name}_{period}.{extension}"), + "legacy_file": legacy_file, } - def _download_single_file( - self, task: Dict[str, Any], force: bool = False, dry_run: bool = True + def _run_tasks( + self, tasks: List[Dict[str, Any]], force: bool, dry_run: bool ) -> DownloadResult: """ - Download a single file based on a task dictionary. + Download all tasks, with up to self.max_workers CDS requests in parallel. + """ - Parameters - ---------- - task : Dict[str, Any] - Task dictionary containing download information (output_file, template, etc.). - force : bool, optional - Force re-download even if file exists. Default is False. - dry_run : bool, optional - If True, only check what would be downloaded. Default is True. + result = self.create_download_result() + if not tasks: + return self.finalize_download_result( + result, "No valid download tasks found" + ) - Returns - ------- - DownloadResult - Result with information about the downloaded, skipped, or error file. + self.logger.info( + f"Prepared {len(tasks)} download tasks ({self.max_workers} in parallel)" + ) + + with ThreadPoolExecutor(max_workers=self.max_workers) as executor: + futures = [ + executor.submit(self._run_task, task, force, dry_run) for task in tasks + ] + for future in futures: + result.merge(future.result()) + + return self.finalize_download_result(result) + + def _run_task( + self, task: Dict[str, Any], force: bool, dry_run: bool + ) -> DownloadResult: + """ + Download a single task. """ result = DownloadResult() - output_file = task["output_file"] - variable = task["variable"] - variable_config = task["variable_config"] - template = task["template"] - - if not dry_run: - os.makedirs(os.path.dirname(output_file), exist_ok=True) - - try: - # Check if file already exists - if not force and os.path.exists(output_file): - if dry_run: - result.add_skipped(output_file, "File already exists (dry run)") - else: - result.add_downloaded(output_file) - return result + legacy_file = task["legacy_file"] + if not force and legacy_file and os.path.exists(legacy_file): + result.add_skipped(legacy_file, "Already downloaded (legacy yearly file)") + return result - if dry_run: - result.add_skipped(output_file, f"Would download {variable} (dry run)") - return result - - # Download file - self.logger.debug(f"Downloading: {variable} to {output_file}") + def fetch(part_file: str) -> None: self.client.retrieve( - name=variable_config["dataset"], - request=template, - target=output_file, + name=task["dataset"], request=task["request"], target=part_file ) - result.add_downloaded(output_file) - self.logger.info(f"Downloaded: {output_file}") - except Exception as e: - self.logger.error(f"Error downloading {output_file}: {e}") - result.add_error(output_file, e) + self.fetch_file( + task["output_file"], fetch, result, force=force, dry_run=dry_run + ) return result diff --git a/bluemath_tk/downloaders/csiro/CAWCR/CAWCR_config.json b/bluemath_tk/downloaders/csiro/CAWCR/CAWCR_config.json new file mode 100644 index 0000000..afa4751 --- /dev/null +++ b/bluemath_tk/downloaders/csiro/CAWCR/CAWCR_config.json @@ -0,0 +1,51 @@ +{ + "description": "CAWCR Wave Hindcast, WAVEWATCH III forced by CFSR / CFSv2", + "fileserver_url": "https://data-cbr.csiro.au/thredds5/fileServer/catch_all/", + "opendap_url": "https://data-cbr.csiro.au/thredds5/dodsC/catch_all/", + "catalog_url": "https://data-cbr.csiro.au/thredds5/catalog/catch_all/", + "datasets": { + "gridded": { + "description": "Global 24 arc-minute grid, hourly, all variables in one file per month", + "remote_dir": "CMAR_CAWCR-Wave_archive/CAWCR_Wave_Hindcast_aggregate/gridded/", + "filename": "ww3.glob_24m.{year}{month:02d}.nc", + "local_dir": "CAWCR/CSIRO_partition", + "start": "1979-01", + "end": null, + "file_per_variable": false + }, + "spec": { + "description": "Spectra at output stations, hourly, one file per month. The station list and variable names change in 2013-06.", + "remote_dir": "CMAR_CAWCR-Wave_archive/CAWCR_Wave_Hindcast_aggregate/spec/", + "filename": "ww3.{year}{month:02d}_spec.nc", + "local_dir": "CAWCR/CSIRO_spectra/{era}", + "eras": { + "1979": [ + "1979-01", + "2013-05" + ], + "2013": [ + "2013-06", + null + ] + }, + "start": "1979-01", + "end": null, + "file_per_variable": false, + "point_dim": "station", + "drop_variables": [ + "string16", + "station_name", + "frequency1", + "frequency2" + ], + "rename": { + "Efth": "efth", + "depth": "dpt", + "u10m": "wnd", + "udir": "wnddir", + "curr": "cur", + "currdir": "curdir" + } + } + } +} diff --git a/bluemath_tk/downloaders/csiro/WHACS/WHACS_config.json b/bluemath_tk/downloaders/csiro/WHACS/WHACS_config.json new file mode 100644 index 0000000..256f94a --- /dev/null +++ b/bluemath_tk/downloaders/csiro/WHACS/WHACS_config.json @@ -0,0 +1,59 @@ +{ + "description": "Wave Hindcast for the Australian Climate Service (WHACS), WAVEWATCH III v6.07 forced by ERA5", + "fileserver_url": "https://data-cbr.csiro.au/thredds5/fileServer/catch_all/", + "opendap_url": "https://data-cbr.csiro.au/thredds5/dodsC/catch_all/", + "catalog_url": "https://data-cbr.csiro.au/thredds5/catalog/catch_all/", + "datasets": { + "global-smc": { + "description": "Global SMC grid, hourly, one file per variable and month", + "remote_dir": "ACS_WP3_WHACS/ACS_hindcast_DRS/gridded/release/WP3/WHACS/BoM-CSIRO/hindcast/ERA5/ERA5/WHACS/WWIII-v6.07/global-smc/1hr/{variable}/", + "filename": "{variable}_WHACS_hindcast_WHACS_ERA5_1hr_{year}{month:02d}010000-{year}{month:02d}{last_day:02d}2300.nc", + "local_dir": "WHACS/global-smc/{variable}", + "start": "1979-01", + "end": "2023-12", + "file_per_variable": true, + "point_dim": "seapoint", + "variables": { + "hs": "Significant wave height", + "t01": "Mean wave period (m0/m1)", + "t02": "Mean wave period (m0/m2)", + "t0m1": "Mean wave period (m-1/m0)", + "fp": "Peak frequency", + "dir": "Mean wave direction", + "dp": "Peak direction", + "spr": "Directional spread", + "cge": "Wave energy flux", + "uwnd": "Eastward wind", + "vwnd": "Northward wind", + "phs0": "Partition 0 (wind sea) significant wave height", + "phs1": "Partition 1 significant wave height", + "phs2": "Partition 2 significant wave height", + "phs3": "Partition 3 significant wave height", + "ptp0": "Partition 0 (wind sea) peak period", + "ptp1": "Partition 1 peak period", + "ptp2": "Partition 2 peak period", + "ptp3": "Partition 3 peak period", + "pdir0": "Partition 0 (wind sea) mean direction", + "pdir1": "Partition 1 mean direction", + "pdir2": "Partition 2 mean direction", + "pdir3": "Partition 3 mean direction", + "pspr0": "Partition 0 (wind sea) directional spread", + "pspr1": "Partition 1 directional spread", + "pspr2": "Partition 2 directional spread", + "pspr3": "Partition 3 directional spread", + "pdp0": "Partition 0 (wind sea) peak direction", + "pdp1": "Partition 1 peak direction", + "pdp2": "Partition 2 peak direction", + "pdp3": "Partition 3 peak direction", + "ppe0": "Partition 0 (wind sea) peakedness", + "ppe1": "Partition 1 peakedness", + "ppe2": "Partition 2 peakedness", + "ppe3": "Partition 3 peakedness", + "pqp0": "Partition 0 (wind sea) Goda peakedness", + "pqp1": "Partition 1 Goda peakedness", + "pqp2": "Partition 2 Goda peakedness", + "pqp3": "Partition 3 Goda peakedness" + } + } + } +} diff --git a/bluemath_tk/downloaders/csiro/__init__.py b/bluemath_tk/downloaders/csiro/__init__.py new file mode 100644 index 0000000..e142f6a --- /dev/null +++ b/bluemath_tk/downloaders/csiro/__init__.py @@ -0,0 +1,11 @@ +""" +Project: BlueMath_tk +Sub-Module: downloaders.csiro +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) +""" + +from .csiro_downloader import CSIRODownloader + +__all__ = ["CSIRODownloader"] diff --git a/bluemath_tk/downloaders/csiro/csiro_downloader.py b/bluemath_tk/downloaders/csiro/csiro_downloader.py new file mode 100644 index 0000000..4dd5a22 --- /dev/null +++ b/bluemath_tk/downloaders/csiro/csiro_downloader.py @@ -0,0 +1,733 @@ +""" +Downloader for the CSIRO wave hindcasts, served by the CSIRO THREDDS server. + +Products (configs in ``csiro//_config.json``): + +- WHACS: WAVEWATCH III hindcast forced by ERA5, global SMC grid, 1979-2023. +- CAWCR: WAVEWATCH III hindcast, global gridded fields and station spectra, + 1979-present. + +Full files are downloaded through the THREDDS fileServer and time series at +given locations are extracted through OPeNDAP. +""" + +import calendar +import os +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Dict, List, Optional, Sequence, Tuple, Union + +import numpy as np +import pandas as pd +import requests +import xarray as xr + +from .._base_downloaders import BaseDownloader, load_config +from .._download_result import DownloadResult +from .._transfer import ( + PART_SUFFIX, + RemoteFileNotFound, + http_download, + make_session, + remote_size, +) + +_HERE = os.path.dirname(__file__) +EARTH_RADIUS_KM = 6371.0 +MAX_DISTANCE_WARNING_KM = 25.0 + + +def nearest_point_index( + lons: np.ndarray, lats: np.ndarray, lon: float, lat: float +) -> Tuple[int, float]: + """ + Index of the point closest to (lon, lat), using the haversine distance. + + Longitudes may use any convention (-180/180 or 0/360). + + Parameters + ---------- + lons, lats : np.ndarray + Longitudes and latitudes of the candidate points, in degrees. + lon, lat : float + Target longitude and latitude, in degrees. + + Returns + ------- + Tuple[int, float] + Flat index of the nearest point and its distance in km. + """ + + lons = np.radians(np.asarray(lons, dtype=float)).ravel() + lats = np.radians(np.asarray(lats, dtype=float)).ravel() + lon0, lat0 = np.radians(lon), np.radians(lat) + a = ( + np.sin((lats - lat0) / 2) ** 2 + + np.cos(lat0) * np.cos(lats) * np.sin((lons - lon0) / 2) ** 2 + ) + idx = int(np.nanargmin(a)) + distance = 2 * EARTH_RADIUS_KM * np.arcsin(np.sqrt(a[idx])) + + return idx, float(distance) + + +class CSIRODownloader(BaseDownloader): + """ + Downloader for the CSIRO wave hindcasts served by the CSIRO THREDDS server. + + Two ways of getting data are available: + + - ``download_data``: full monthly files through the THREDDS fileServer. + Transfers are resumable, checked against the remote file size and run in + parallel (``max_workers``). + - ``download_point``: time series at one or more locations through OPeNDAP, + without downloading the full files. One file is written per month and, + optionally, all months are combined into a single file. + + Files are saved to base_path_to_download/, so with + base_path_to_download="/lustre/geocean/DATA/CSIRO" the existing GeoOcean + layout is reused and already downloaded files are skipped. + + Examples + -------- + >>> downloader = CSIRODownloader( + ... product="WHACS", + ... base_path_to_download="/lustre/geocean/DATA/CSIRO", + ... ) + >>> result = downloader.download_data( + ... dataset="global-smc", + ... variables=["phs1", "ptp1"], + ... start="2000-01", + ... end="2000-12", + ... dry_run=False, + ... ) + >>> result = downloader.download_point( + ... dataset="global-smc", + ... lon=-4.0, + ... lat=44.0, + ... variables=["hs", "t02"], + ... dry_run=False, + ... ) + """ + + products_configs = { + "WHACS": load_config(_HERE, "WHACS", "WHACS_config.json"), + "CAWCR": load_config(_HERE, "CAWCR", "CAWCR_config.json"), + } + + def __init__( + self, + product: str, + base_path_to_download: str, + max_workers: int = 2, + debug: bool = True, + **kwargs, + ) -> None: + """ + Initialize the CSIRODownloader. + + Parameters + ---------- + product : str + The product to download data from ("WHACS" or "CAWCR"). + base_path_to_download : str + Base path where downloaded files will be stored. + max_workers : int, optional + Number of files downloaded in parallel by download_data. Default is 2. + debug : bool, optional + If True, sets logger to DEBUG level. Default is True. + **kwargs + Passed to BaseDownloader (e.g., retries, retry_backoff). + + Raises + ------ + ValueError + If the product configuration is not found. + """ + + super().__init__( + product=product, + base_path_to_download=base_path_to_download, + debug=debug, + **kwargs, + ) + + self.max_workers = max_workers + self._local = threading.local() + self._grid_cache: Dict[tuple, Tuple[np.ndarray, np.ndarray]] = {} + self._exclude_attributes.extend(["_local", "_grid_cache"]) + + self.set_logger_name( + f"CSIRODownloader-{product}", level="DEBUG" if debug else "INFO" + ) + self.logger.info(f"---- CSIRO DOWNLOADER INITIALIZED ({product}) ----") + + @property + def session(self) -> requests.Session: + """ + HTTP session of the current thread. + + Returns + ------- + requests.Session + Session with retries for connection errors and 5xx responses. + """ + if not hasattr(self._local, "session"): + self._local.session = make_session() + return self._local.session + + def list_variables(self, dataset: str) -> List[str]: + """ + List the variables of a dataset with one file per variable. + + Parameters + ---------- + dataset : str + The dataset name. + + Returns + ------- + List[str] + List of variable names (empty if all variables are in the same file). + """ + + return list(self._get_dataset_config(dataset).get("variables", {}).keys()) + + def get_months( + self, dataset: str, start: str = None, end: str = None + ) -> pd.PeriodIndex: + """ + Months to process, limited to the dataset availability. + + Parameters + ---------- + dataset : str + The dataset name. + start : str, optional + First month ("YYYY-MM"). Default is the dataset start. + end : str, optional + Last month ("YYYY-MM"). Default is the dataset end or, for datasets + that are still updated, the current month. + + Returns + ------- + pd.PeriodIndex + Monthly periods. + """ + + dataset_config = self._get_dataset_config(dataset) + first = pd.Period(dataset_config["start"], freq="M") + last = ( + pd.Period(dataset_config["end"], freq="M") + if dataset_config.get("end") + else pd.Period.now(freq="M") + ) + start = max(pd.Period(start, freq="M"), first) if start else first + end = min(pd.Period(end, freq="M"), last) if end else last + + return pd.period_range(start, end, freq="M") + + def get_url( + self, + dataset: str, + period: pd.Period, + variable: str = None, + service: str = "fileserver", + ) -> str: + """ + URL of a monthly file. + + Parameters + ---------- + dataset : str + The dataset name. + period : pd.Period + The month. + variable : str, optional + The variable, for datasets with one file per variable. + service : str, optional + "fileserver" (download) or "opendap" (remote access). Default is "fileserver". + + Returns + ------- + str + The URL. + """ + + dataset_config = self._get_dataset_config(dataset) + base_url = self.product_config[f"{service}_url"] + template_args = self._template_args(period, variable) + remote_dir = dataset_config["remote_dir"].format(**template_args) + filename = dataset_config["filename"].format(**template_args) + + return f"{base_url}{remote_dir}{filename}" + + def get_local_path( + self, dataset: str, period: pd.Period, variable: str = None + ) -> str: + """ + Local path of a monthly file. + + Parameters + ---------- + dataset : str + The dataset name. + period : pd.Period + The month. + variable : str, optional + The variable, for datasets with one file per variable. + + Returns + ------- + str + The local path. + """ + + dataset_config = self._get_dataset_config(dataset) + template_args = self._template_args(period, variable) + template_args["era"] = self._get_era(dataset_config, period) + local_dir = dataset_config["local_dir"].format(**template_args) + filename = dataset_config["filename"].format(**template_args) + + # local_dir is written with "/" in the config: use the native separator + return os.path.normpath( + os.path.join(self.base_path_to_download, local_dir, filename) + ) + + def download_data( + self, + dataset: str, + variables: List[str] = None, + start: str = None, + end: str = None, + force: bool = False, + verify_existing: bool = True, + dry_run: bool = True, + ) -> DownloadResult: + """ + Download full monthly files. + + Months that are not (yet) on the server are reported as skipped, so the + same call can be repeated to keep a dataset up to date. + + Parameters + ---------- + dataset : str + The dataset to download. Use list_datasets() to see available datasets. + variables : List[str], optional + Variables to download, required for datasets with one file per + variable (see list_variables). Ignored otherwise. + start : str, optional + First month ("YYYY-MM"). Default is the dataset start. + end : str, optional + Last month ("YYYY-MM"). Default is the last available month. + force : bool, optional + Download again even if the file exists. Default is False. + verify_existing : bool, optional + Compare the size of existing files with the remote size and download + them again if they differ (e.g. interrupted transfers). Default is True. + dry_run : bool, optional + If True, only check what would be downloaded. Default is True. + + Returns + ------- + DownloadResult + Result with information about downloaded, skipped, and error files. + The remote size of each checked file is in metadata["remote_sizes"]. + """ + + dataset_config = self._get_dataset_config(dataset) + file_variables = self._check_variables(dataset, dataset_config, variables) + + tasks = [ + ( + self.get_url(dataset, period, variable), + self.get_local_path(dataset, period, variable), + ) + for variable in file_variables + for period in self.get_months(dataset, start, end) + ] + self.logger.info(f"Checking {len(tasks)} files of {self.product}/{dataset}") + + result = self.create_download_result() + with ThreadPoolExecutor(max_workers=self.max_workers) as executor: + futures = [ + executor.submit( + self._download_file, url, path, force, verify_existing, dry_run + ) + for url, path in tasks + ] + for future in futures: + result.merge(future.result()) + + return self.finalize_download_result(result) + + def download_point( + self, + dataset: str, + lon: Union[float, Sequence[float]], + lat: Union[float, Sequence[float]], + variables: List[str] = None, + start: str = None, + end: str = None, + output_dir: str = None, + name: str = None, + combine: bool = True, + force: bool = False, + dry_run: bool = True, + ) -> DownloadResult: + """ + Extract time series at the nearest model point(s) through OPeNDAP. + + One file per month is saved to output_dir/{name}_{YYYYMM}.nc. If combine + is True, all available months are also merged into + output_dir/{name}_{dataset}_{start}_{end}.nc. + + Parameters + ---------- + dataset : str + The dataset to extract from. + lon, lat : float or Sequence[float] + Target location(s), in degrees. + variables : List[str], optional + Variables to extract. Required for datasets with one file per variable. + For other datasets, None extracts all variables. + start : str, optional + First month ("YYYY-MM"). Default is the dataset start. + end : str, optional + Last month ("YYYY-MM"). Default is the last available month. + output_dir : str, optional + Output directory. Default is + base_path_to_download/product/dataset/points/name. + name : str, optional + Name of the extraction, used for file names. Required for several + locations. Default is "lon{lon}_lat{lat}". + combine : bool, optional + Merge all monthly files into a single file. Default is True. + force : bool, optional + Extract again even if the monthly file exists. Default is False. + dry_run : bool, optional + If True, only check what would be extracted. Default is True. + + Returns + ------- + DownloadResult + Result with the monthly files. The combined file, if any, is in + result.metadata["combined_file"]. + """ + + dataset_config = self._get_dataset_config(dataset) + lons = np.atleast_1d(np.asarray(lon, dtype=float)) + lats = np.atleast_1d(np.asarray(lat, dtype=float)) + if lons.shape != lats.shape: + raise ValueError("lon and lat must have the same length") + if dataset_config["file_per_variable"]: + self._check_variables(dataset, dataset_config, variables) + if name is None: + if len(lons) > 1: + raise ValueError("name is required when extracting several points") + name = f"lon{lons[0]:.3f}_lat{lats[0]:.3f}" + if output_dir is None: + output_dir = os.path.join( + self.base_path_to_download, self.product, dataset, "points", name + ) + + months = self.get_months(dataset, start, end) + self.logger.info( + f"Extracting {len(months)} months of {self.product}/{dataset} at {name}" + ) + + result = self.create_download_result() + monthly_files = [] + for period in months: + output_file = os.path.join( + output_dir, f"{name}_{period.year}{period.month:02d}.nc" + ) + monthly_files.append(output_file) + + def fetch(part_file: str, period: pd.Period = period) -> None: + point_ds = self._extract_point_month( + dataset, dataset_config, period, variables, lons, lats + ) + point_ds.to_netcdf(part_file) + + self.fetch_file(output_file, fetch, result, force=force, dry_run=dry_run) + + available_files = [f for f in monthly_files if os.path.exists(f)] + if combine and not dry_run and available_files: + combined_file = os.path.join( + output_dir, + f"{name}_{dataset}_{months[0].strftime('%Y%m')}_{months[-1].strftime('%Y%m')}.nc", + ) + try: + self.combine_point_files(available_files, combined_file) + result.metadata["combined_file"] = combined_file + self.logger.info(f"Combined file: {combined_file}") + except Exception as e: + self.logger.error(f"Error combining point files: {e}") + result.add_error(combined_file, e) + + return self.finalize_download_result(result) + + @staticmethod + def combine_point_files(files: List[str], output_file: str) -> None: + """ + Concatenate monthly point files along time into a single NetCDF file. + + Parameters + ---------- + files : List[str] + Monthly files written by download_point. + output_file : str + Path of the combined file. It is overwritten if it exists. + """ + + part_file = output_file + PART_SUFFIX + with xr.open_mfdataset( + sorted(files), + combine="nested", + concat_dim="time", + data_vars="minimal", + coords="different", + ) as ds: + ds.drop_duplicates("time").to_netcdf(part_file) + os.replace(part_file, output_file) + + def _get_dataset_config(self, dataset: str) -> dict: + if dataset not in self.list_datasets(): + raise ValueError( + f"Dataset '{dataset}' not found. Available: {self.list_datasets()}" + ) + return self.product_config["datasets"][dataset] + + def _check_variables( + self, dataset: str, dataset_config: dict, variables: Optional[List[str]] + ) -> List[Optional[str]]: + """ + Variables that define the files to process ([None] if one file holds all). + """ + + if not dataset_config["file_per_variable"]: + return [None] + if not variables: + raise ValueError( + f"variables is required for {self.product}/{dataset}. " + f"Available: {self.list_variables(dataset)}" + ) + unknown = set(variables) - set(self.list_variables(dataset)) + if unknown: + raise ValueError( + f"Unknown variables {sorted(unknown)} for {self.product}/{dataset}. " + f"Available: {self.list_variables(dataset)}" + ) + return list(variables) + + @staticmethod + def _template_args(period: pd.Period, variable: Optional[str]) -> dict: + return { + "year": period.year, + "month": period.month, + "last_day": calendar.monthrange(period.year, period.month)[1], + "variable": variable, + } + + @staticmethod + def _get_era(dataset_config: dict, period: pd.Period) -> Optional[str]: + """ + Name of the era (file format period) a month belongs to, if any. + """ + + for era, (era_start, era_end) in dataset_config.get("eras", {}).items(): + if period >= pd.Period(era_start, freq="M") and ( + era_end is None or period <= pd.Period(era_end, freq="M") + ): + return era + return None + + def _download_file( + self, + url: str, + output_file: str, + force: bool, + verify_existing: bool, + dry_run: bool, + ) -> DownloadResult: + """ + Download a single file, checking the remote size first. + """ + + result = DownloadResult() + exists = os.path.exists(output_file) + if exists and not force and not verify_existing: + result.add_skipped(output_file, "Already downloaded") + return result + + try: + size = remote_size(self.session, url) + except RemoteFileNotFound: + reason = "Not available on server" + result.add_skipped( + output_file, reason + (" (local copy kept)" if exists else "") + ) + return result + except Exception as e: + self.logger.error(f"Error checking {url}: {e}") + result.add_error(output_file, e) + return result + + if size is not None: + result.metadata.setdefault("remote_sizes", {})[output_file] = size + + if exists and not force: + local_size = os.path.getsize(output_file) + if size is None or local_size == size: + result.add_skipped(output_file, "Already downloaded") + return result + message = f"Local size {local_size} differs from remote size {size}" + if dry_run: + result.add_skipped( + output_file, f"{message}, would download again (dry run)" + ) + return result + self.logger.warning(f"{message}, downloading again: {output_file}") + force = True + + def fetch(part_file: str) -> None: + http_download(self.session, url, part_file, expected_size=size, resume=True) + + self.fetch_file( + output_file, fetch, result, force=force, dry_run=dry_run, resumable=True + ) + + return result + + def _extract_point_month( + self, + dataset: str, + dataset_config: dict, + period: pd.Period, + variables: Optional[List[str]], + lons: np.ndarray, + lats: np.ndarray, + ) -> xr.Dataset: + """ + Extract one month at the nearest point(s) through OPeNDAP. + """ + + rename = dataset_config.get("rename", {}) + file_variables = variables if dataset_config["file_per_variable"] else [None] + + pieces = [] + for variable in file_variables: + # OPeNDAP errors for missing files are not explicit, so check first + remote_size(self.session, self.get_url(dataset, period, variable)) + url = self.get_url(dataset, period, variable, service="opendap") + with xr.open_dataset( + url, drop_variables=dataset_config.get("drop_variables") + ) as ds: + ds = ds.rename({k: v for k, v in rename.items() if k in ds.variables}) + ds = ds.set_coords( + [c for c in ("longitude", "latitude") if c in ds.data_vars] + ) + point_ds = self._select_points(dataset, dataset_config, ds, lons, lats) + if variable is not None: + point_ds = point_ds[[variable]] + elif variables: + missing = set(variables) - set(point_ds.data_vars) + if missing: + raise ValueError(f"Variables {sorted(missing)} not in {url}") + point_ds = point_ds[variables] + pieces.append(point_ds.load()) + + point_ds = xr.merge(pieces, compat="override", combine_attrs="override") + # Monthly files also hold the first time step of the next month + in_period = point_ds.time.to_index().to_period("M") == period + point_ds = point_ds.isel(time=np.flatnonzero(in_period)) + point_ds.attrs.update( + { + "source": f"CSIRO {self.product} {dataset}", + "target_longitude": lons.tolist(), + "target_latitude": lats.tolist(), + } + ) + + return point_ds + + def _select_points( + self, + dataset: str, + dataset_config: dict, + ds: xr.Dataset, + lons: np.ndarray, + lats: np.ndarray, + ) -> xr.Dataset: + """ + Lazily select the nearest point(s) of a dataset. + + Datasets with a point dimension (stations, SMC seapoints) are indexed + along it. Regular grids are indexed along latitude and longitude. + """ + + point_dim = dataset_config.get("point_dim") + lon_values, lat_values = self._get_grid(dataset, ds, point_dim) + + nearest = [ + nearest_point_index(lon_values, lat_values, lon, lat) + for lon, lat in zip(lons, lats) + ] + distances = [distance for _, distance in nearest] + for (lon, lat), distance in zip(zip(lons, lats), distances): + message = f"Nearest point to ({lon}, {lat}) is {distance:.1f} km away" + if distance > MAX_DISTANCE_WARNING_KM: + self.logger.warning(message) + else: + self.logger.debug(message) + + if point_dim is not None: + indices = np.array([idx for idx, _ in nearest]) + if len(indices) == 1: + selected = ds.isel({point_dim: int(indices[0])}) + else: + unique, inverse = np.unique(indices, return_inverse=True) + selected = ( + ds.isel({point_dim: unique}).load().isel({point_dim: inverse}) + ) + else: + lat_idx, lon_idx = np.unravel_index( + [idx for idx, _ in nearest], lon_values.shape + ) + if len(lat_idx) == 1: + selected = ds.isel(latitude=int(lat_idx[0]), longitude=int(lon_idx[0])) + else: + selected = ds.isel( + latitude=xr.DataArray(lat_idx, dims="point"), + longitude=xr.DataArray(lon_idx, dims="point"), + ) + + return selected.assign_attrs(nearest_point_distance_km=distances) + + def _get_grid( + self, dataset: str, ds: xr.Dataset, point_dim: Optional[str] + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Longitudes and latitudes of all points, cached when they do not change. + """ + + longitude, latitude = ds["longitude"], ds["latitude"] + + # Station positions stored per time step: read the first one, no caching + if "time" in longitude.dims: + return ( + longitude.isel(time=0).values, + latitude.isel(time=0).values, + ) + + key = (dataset, longitude.size, latitude.size) + if key not in self._grid_cache: + if point_dim is None: + lon_grid, lat_grid = np.meshgrid(longitude.values, latitude.values) + self._grid_cache[key] = (lon_grid, lat_grid) + else: + self._grid_cache[key] = (longitude.values, latitude.values) + + return self._grid_cache[key] diff --git a/bluemath_tk/downloaders/ecmwf/__init__.py b/bluemath_tk/downloaders/ecmwf/__init__.py index e69de29..46890b8 100644 --- a/bluemath_tk/downloaders/ecmwf/__init__.py +++ b/bluemath_tk/downloaders/ecmwf/__init__.py @@ -0,0 +1,11 @@ +""" +Project: BlueMath_tk +Sub-Module: downloaders.ecmwf +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) +""" + +from .ecmwf_downloader import ECMWFDownloader + +__all__ = ["ECMWFDownloader"] diff --git a/bluemath_tk/downloaders/ecmwf/ecmwf_downloader.py b/bluemath_tk/downloaders/ecmwf/ecmwf_downloader.py index f102264..7d768b3 100644 --- a/bluemath_tk/downloaders/ecmwf/ecmwf_downloader.py +++ b/bluemath_tk/downloaders/ecmwf/ecmwf_downloader.py @@ -1,15 +1,32 @@ -import json +""" +Downloader for ECMWF Open Data forecasts, through ``ecmwf-opendata``. + +Products (configs in ``ecmwf//_config.json``): + +- OpenData: IFS and AIFS real-time forecasts. +""" + import os +from datetime import date, datetime, timedelta +from typing import Any, Dict +import pandas as pd from ecmwf.opendata import Client -from .._base_downloaders import BaseDownloader +from .._base_downloaders import BaseDownloader, load_config from .._download_result import DownloadResult +_HERE = os.path.dirname(__file__) + class ECMWFDownloader(BaseDownloader): """ - This is the main class to download data from the ECMWF. + Downloader for ECMWF Open Data real-time forecasts. + + Each call downloads one GRIB2 file for one forecast run, saved to + base_path_to_download/product/dataset/model/resolution// + __.grib2, so files of different runs never overwrite + each other. Without ``date``, the latest available run is used. Examples -------- @@ -17,7 +34,8 @@ class ECMWFDownloader(BaseDownloader): ... product="OpenData", ... base_path_to_download="./ecmwf_data", ... model="ifs", - ... resolution="0p25" + ... resolution="0p25", + ... source="aws", ... ) >>> result = downloader.download_data( ... dataset="forecast_data", @@ -30,13 +48,7 @@ class ECMWFDownloader(BaseDownloader): """ products_configs = { - "OpenData": json.load( - open( - os.path.join( - os.path.dirname(__file__), "OpenData", "OpenData_config.json" - ) - ) - ) + "OpenData": load_config(_HERE, "OpenData", "OpenData_config.json"), } def __init__( @@ -45,7 +57,9 @@ def __init__( base_path_to_download: str, model: str = "ifs", resolution: str = "0p25", + source: str = "ecmwf", debug: bool = True, + **kwargs, ) -> None: """ Initialize the ECMWFDownloader. @@ -60,8 +74,14 @@ def __init__( The model to download data from (e.g., "ifs", "aifs"). Default is "ifs". resolution : str, optional The resolution to download data from (e.g., "0p25"). Default is "0p25". + source : str, optional + Where to download from: "ecmwf" (the ECMWF portal, limited to 500 + simultaneous connections) or one of its mirrors "aws", "azure", + "google". Default is "ecmwf". debug : bool, optional If True, sets logger to DEBUG level. Default is True. + **kwargs + Passed to BaseDownloader (e.g., retries, retry_backoff). Raises ------ @@ -70,16 +90,12 @@ def __init__( """ super().__init__( - product=product, base_path_to_download=base_path_to_download, debug=debug + product=product, + base_path_to_download=base_path_to_download, + debug=debug, + **kwargs, ) - self._product_config = self.products_configs.get(product) - if self._product_config is None: - available_products = list(self.products_configs.keys()) - raise ValueError( - f"{product} configuration not found. Available products: {available_products}" - ) - self.set_logger_name( f"ECMWFDownloader-{product}", level="DEBUG" if debug else "INFO" ) @@ -94,9 +110,8 @@ def __init__( raise ValueError( f"Resolution {resolution} not supported for {self.product}" ) - # Always initialize client (will skip API calls in dry_run mode) self._client = Client( - source="ecmwf", + source=source, model=model, resol=resolution, preserve_request_order=False, @@ -104,21 +119,9 @@ def __init__( ) self.logger.info(f"---- ECMWF DOWNLOADER INITIALIZED ({product}) ----") - # Set the model and resolution parameters self.model = model self.resolution = resolution - - @property - def product_config(self) -> dict: - """ - Product configuration dictionary loaded from config file. - - Returns - ------- - dict - Product configuration dictionary. - """ - return self._product_config + self.source = source @property def client(self) -> Client: @@ -132,43 +135,6 @@ def client(self) -> Client: """ return self._client - def download_data( - self, - dry_run: bool = True, - *args, - **kwargs, - ) -> DownloadResult: - """ - Download data for the product. - - Routes to product-specific download methods based on the product type. - - Parameters - ---------- - dry_run : bool, optional - If True, only check what would be downloaded without actually downloading. - Default is True. - *args - Arguments passed to product-specific download method. - **kwargs - Keyword arguments passed to product-specific download method. - - Returns - ------- - DownloadResult - Result with information about downloaded, skipped, and error files. - - Raises - ------ - ValueError - If the product is not supported. - """ - - if self.product == "OpenData": - return self.download_data_open_data(dry_run=dry_run, *args, **kwargs) - else: - raise ValueError(f"Download for product {self.product} not supported") - def download_data_open_data( self, dataset: str, @@ -179,8 +145,10 @@ def download_data_open_data( """ Download data for the OpenData product. - Downloads files based on the specified parameters. Files are saved to: - base_path_to_download/product/dataset/model/resolution/filename.grib2 + Downloads one GRIB2 file with the requested parameters and steps of one + forecast run. Files are saved to: + base_path_to_download/product/dataset/model/resolution// + __.grib2 Parameters ---------- @@ -193,7 +161,10 @@ def download_data_open_data( If True, only check what would be downloaded. Default is True. **kwargs Keyword arguments passed to the ECMWF client retrieve method - (e.g., param, step, type). + (e.g., param, step, type). ``date`` (YYYYMMDD, a datetime, or 0, -1, + ... for days before today) and ``time`` (run hour) select the + forecast run; without ``date`` the latest run is used. Finding the + latest run needs a request to the server, also in dry runs. Returns ------- @@ -214,64 +185,66 @@ def download_data_open_data( result = self.create_download_result() - try: - # Extract parameters from kwargs - if "param" in kwargs: - variables = kwargs["param"] - else: - variables = [] - if "step" in kwargs: - steps = kwargs["step"] - if not isinstance(steps, list): - steps = [steps] - else: - steps = [] - if "type" in kwargs: - type = kwargs["type"] - else: - type = "fc" - - # Construct output file path: base_path/product/dataset/model/resolution/filename.grib2 - output_grib_file = os.path.join( - self.base_path_to_download, - self.product, - dataset, - self.model, - self.resolution, - f"{'_'.join(variables)}_{'_'.join(str(step) for step in steps)}_{type}.grib2", - ) + # Pin the forecast run, so the file always matches its path + run = self._get_run_datetime(kwargs) + kwargs["date"] = run.strftime("%Y%m%d") + kwargs["time"] = run.hour + + variables = kwargs.get("param", []) + if not isinstance(variables, list): + variables = [variables] + steps = kwargs.get("step", []) + if not isinstance(steps, list): + steps = [steps] + type = kwargs.get("type", "fc") + + output_grib_file = os.path.join( + self.base_path_to_download, + self.product, + dataset, + self.model, + self.resolution, + run.strftime("%Y%m%d%H"), + f"{'_'.join(variables)}_{'_'.join(str(step) for step in steps)}_{type}.grib2", + ) - # Skip if file already exists (unless force=True) - if not force and os.path.exists(output_grib_file): - result.add_skipped(output_grib_file, "Already downloaded") - return self.finalize_download_result(result) + def fetch(part_file: str) -> None: + self.client.retrieve(target=part_file, **kwargs) - # Handle dry run: record as skipped without actual download - if dry_run: - result.add_skipped(output_grib_file, "Would download (dry run)") - return self.finalize_download_result(result) + self.fetch_file(output_grib_file, fetch, result, force=force, dry_run=dry_run) - # Attempt to download the file - try: - # Create local directory structure if needed - os.makedirs(os.path.dirname(output_grib_file), exist_ok=True) + return self.finalize_download_result(result) - # Download the file - self.logger.debug(f"Downloading: {output_grib_file}") - self.client.retrieve( - target=output_grib_file, - **kwargs, - ) + def _get_run_datetime(self, request: Dict[str, Any]) -> datetime: + """ + Date and hour of the forecast run selected by a request. - result.add_downloaded(output_grib_file) - self.logger.info(f"Downloaded: {output_grib_file}") + Parameters + ---------- + request : Dict[str, Any] + Retrieve keyword arguments. ``date`` may be missing (latest run), an + integer (0 today, -1 yesterday, ...), "YYYYMMDD" or a date/datetime. + ``time`` is the run hour (default 0). - except Exception as e: - result.add_error(output_grib_file, e) - self.logger.error(f"Error downloading {output_grib_file}: {e}") + Returns + ------- + datetime + The run date and hour. + """ - return self.finalize_download_result(result) + run_date = request.get("date") + if run_date is None: + latest_request = { + key: value + for key, value in request.items() + if key not in ("date", "time") + } + return self.client.latest(**latest_request) + + if isinstance(run_date, int) and run_date <= 0: + run_day = date.today() + timedelta(days=run_date) + else: + run_day = pd.Timestamp(str(run_date)).date() + hour = int(str(request.get("time", 0)).split(":")[0]) - except Exception as e: - result.add_error("download_operation", e) - return self.finalize_download_result(result) + return datetime(run_day.year, run_day.month, run_day.day, hour) diff --git a/bluemath_tk/downloaders/noaa/NDBC/NDBC_config.json b/bluemath_tk/downloaders/noaa/NDBC/NDBC_config.json index 44e2877..343b15d 100644 --- a/bluemath_tk/downloaders/noaa/NDBC/NDBC_config.json +++ b/bluemath_tk/downloaders/noaa/NDBC/NDBC_config.json @@ -13,17 +13,6 @@ "buoy_id", "year" ] - }, - "forecast_data": { - "name": "NOAA GFS Forecast Data", - "description": "GFS 0.25 degree forecast data", - "base_url": "https://nomads.ncep.noaa.gov/dods/gfs_0p25_1hr", - "template": { - "date": null - }, - "mandatory_fields": [ - "date" - ] } }, "data_types": { @@ -95,21 +84,6 @@ "url_pattern": "historical/swden/{buoy_id}w{year}.txt.gz" } } - }, - "wind_forecast": { - "name": "wind_forecast", - "long_name": "GFS Wind Forecast", - "description": "Wind speed and direction forecast from GFS model", - "dataset": "forecast_data", - "file_format": "netcdf", - "variables": [ - "ugrd10m", - "vgrd10m" - ], - "output_variables": { - "u10": "ugrd10m", - "v10": "vgrd10m" - } } } -} \ No newline at end of file +} diff --git a/bluemath_tk/downloaders/noaa/__init__.py b/bluemath_tk/downloaders/noaa/__init__.py index e69de29..b46d1b7 100644 --- a/bluemath_tk/downloaders/noaa/__init__.py +++ b/bluemath_tk/downloaders/noaa/__init__.py @@ -0,0 +1,11 @@ +""" +Project: BlueMath_tk +Sub-Module: downloaders.noaa +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) +""" + +from .noaa_downloader import NOAADownloader + +__all__ = ["NOAADownloader"] diff --git a/bluemath_tk/downloaders/noaa/noaa_downloader.py b/bluemath_tk/downloaders/noaa/noaa_downloader.py index 09709d6..86bbcaf 100644 --- a/bluemath_tk/downloaders/noaa/noaa_downloader.py +++ b/bluemath_tk/downloaders/noaa/noaa_downloader.py @@ -1,18 +1,28 @@ +""" +Downloader and readers for NOAA products. + +Products (configs in ``noaa//_config.json``): + +- NDBC: National Data Buoy Center buoy data (bulk parameters, wave spectra, + directional spectra). + +The ``read_*`` functions read the files written by NOAADownloader. +""" + import gzip import io -import json import os -import shutil from datetime import datetime from pathlib import Path from typing import List, Optional, Tuple, Union import pandas as pd -import requests -import xarray as xr -from .._base_downloaders import BaseDownloader +from .._base_downloaders import BaseDownloader, load_config from .._download_result import DownloadResult +from .._transfer import RemoteFileNotFound, http_download, http_get, make_session + +_HERE = os.path.dirname(__file__) def read_bulk_parameters( @@ -39,35 +49,38 @@ def read_bulk_parameters( if isinstance(years, int): years = [years] - all_data = [] - for year in years: - file_path = os.path.join( - base_path, - "NDBC", - "buoy_data", - buoy_id, - f"buoy_{buoy_id}_bulk_parameters.csv", - ) - try: - df = pd.read_csv(file_path) - df["datetime"] = pd.to_datetime( - df["YYYY"].astype(str) - + "-" - + df["MM"].astype(str).str.zfill(2) - + "-" - + df["DD"].astype(str).str.zfill(2) - + " " - + df["hh"].astype(str).str.zfill(2) - + ":" - + df["mm"].astype(str).str.zfill(2) - ) - all_data.append(df) - except FileNotFoundError: - print(f"No bulk parameters file found for buoy {buoy_id} year {year}") + # All years of a buoy are stored in a single file + file_path = os.path.join( + base_path, + "NDBC", + "buoy_data", + buoy_id, + f"buoy_{buoy_id}_bulk_parameters.csv", + ) + try: + df = pd.read_csv(file_path) + except FileNotFoundError: + print(f"No bulk parameters file found for buoy {buoy_id}") + return None - if all_data: - return pd.concat(all_data, ignore_index=True).sort_values("datetime") - return None + df = df[df["YYYY"].isin([int(year) for year in years])].copy() + if df.empty: + print(f"No bulk parameters found for buoy {buoy_id} years {years}") + return None + + df["datetime"] = pd.to_datetime( + df["YYYY"].astype(str) + + "-" + + df["MM"].astype(str).str.zfill(2) + + "-" + + df["DD"].astype(str).str.zfill(2) + + " " + + df["hh"].astype(str).str.zfill(2) + + ":" + + df["mm"].astype(str).str.zfill(2) + ) + + return df.sort_values("datetime").reset_index(drop=True) def read_wave_spectra( @@ -295,7 +308,10 @@ def read_directional_spectra( class NOAADownloader(BaseDownloader): """ - This is the main class to download data from NOAA. + Downloader for NOAA NDBC buoy data. + + Buoy data are saved per buoy under base_path_to_download/NDBC/buoy_data/ + / and can be read back with the ``read_*`` functions of this module. Examples -------- @@ -314,9 +330,7 @@ class NOAADownloader(BaseDownloader): """ products_configs = { - "NDBC": json.load( - open(os.path.join(os.path.dirname(__file__), "NDBC", "NDBC_config.json")) - ) + "NDBC": load_config(_HERE, "NDBC", "NDBC_config.json"), } def __init__( @@ -324,6 +338,7 @@ def __init__( product: str, base_path_to_download: str, debug: bool = True, + **kwargs, ) -> None: """ Initialize the NOAA downloader. @@ -336,6 +351,8 @@ def __init__( The base path to download the data to. debug : bool, optional Whether to run in debug mode. Default is True. + **kwargs + Passed to BaseDownloader (e.g., retries, retry_backoff). Raises ------ @@ -344,33 +361,19 @@ def __init__( """ super().__init__( - product=product, base_path_to_download=base_path_to_download, debug=debug + product=product, + base_path_to_download=base_path_to_download, + debug=debug, + **kwargs, ) - self._product_config = self.products_configs.get(product) - if self._product_config is None: - available_products = list(self.products_configs.keys()) - raise ValueError( - f"Product '{product}' not found. Available: {available_products}" - ) - self.set_logger_name( f"NOAADownloader-{product}", level="DEBUG" if debug else "INFO" ) + self._session = make_session() + self._exclude_attributes.append("_session") self.logger.info(f"---- NOAA DOWNLOADER INITIALIZED ({product}) ----") - @property - def product_config(self) -> dict: - """ - Product configuration dictionary loaded from config file. - - Returns - ------- - dict - Product configuration dictionary. - """ - return self._product_config - @property def data_types(self) -> dict: """ @@ -394,73 +397,8 @@ def list_data_types(self) -> List[str]: """ return list(self.data_types.keys()) - def _check_file_exists( - self, file_path: str, result: DownloadResult, force: bool, dry_run: bool - ) -> bool: - """ - Check if file exists and handle accordingly. - - Parameters - ---------- - file_path : str - Path to the file to check. - result : DownloadResult - The download result to update. - force : bool - Whether to force re-download. - dry_run : bool - If True, only check files without downloading. - - Returns - ------- - bool - True if should skip download (file exists or dry_run mode), False otherwise. - """ - - if not force and os.path.exists(file_path): - result.add_skipped(file_path, "File already exists") - return True - - if dry_run: - result.add_skipped(file_path, "File does not exist (dry run)") - return True - - return False - - def download_data(self, dry_run: bool = True, *args, **kwargs) -> DownloadResult: - """ - Download data for the product. - - Routes to product-specific download methods based on the product type. - - Parameters - ---------- - dry_run : bool, optional - If True, only check what would be downloaded without actually downloading. - Default is True. - *args - Arguments passed to product-specific download method. - **kwargs - Keyword arguments passed to product-specific download method. - - Returns - ------- - DownloadResult - Result with information about downloaded, skipped, and error files. - - Raises - ------ - ValueError - If the product is not supported. - """ - - if self.product == "NDBC": - return self.download_data_ndbc(dry_run=dry_run, *args, **kwargs) - else: - raise ValueError(f"Download for product {self.product} not supported") - def download_data_ndbc( - self, data_type: str, dry_run: bool = True, **kwargs + self, data_type: str, force: bool = False, dry_run: bool = True, **kwargs ) -> DownloadResult: """ Download data for the NDBC product. @@ -472,17 +410,17 @@ def download_data_ndbc( ---------- data_type : str The data type to download. Available types: + - 'bulk_parameters': Standard meteorological data - 'wave_spectra': Wave spectral density data - 'directional_spectra': Directional wave spectra coefficients - - 'wind_forecast': GFS wind forecast data + force : bool, optional + Force re-download even if the file exists. Default is False. dry_run : bool, optional If True, only check what would be downloaded without actually downloading. Default is True. **kwargs - Additional keyword arguments specific to each data type: - - For bulk_parameters, wave_spectra, directional_spectra: buoy_id, years, force - - For wind_forecast: date, region, force + Keyword arguments of the data type: buoy_id and years. Returns ------- @@ -508,19 +446,15 @@ def download_data_ndbc( if data_type == "bulk_parameters": result = self._download_bulk_parameters( - data_type_config, dataset_config, dry_run=dry_run, **kwargs + data_type_config, dataset_config, force=force, dry_run=dry_run, **kwargs ) elif data_type == "wave_spectra": result = self._download_wave_spectra( - data_type_config, dataset_config, dry_run=dry_run, **kwargs + data_type_config, dataset_config, force=force, dry_run=dry_run, **kwargs ) elif data_type == "directional_spectra": result = self._download_directional_spectra( - data_type_config, dataset_config, dry_run=dry_run, **kwargs - ) - elif data_type == "wind_forecast": - result = self._download_wind_forecast( - data_type_config, dataset_config, dry_run=dry_run, **kwargs + data_type_config, dataset_config, force=force, dry_run=dry_run, **kwargs ) else: raise ValueError(f"Download for data type {data_type} not implemented") @@ -539,6 +473,9 @@ def _download_bulk_parameters( """ Download bulk parameters for a specific buoy and years. + All years are combined into a single file: + base_path_to_download/product/dataset/buoy_id/buoy_{buoy_id}_bulk_parameters.csv + Parameters ---------- data_type_config : dict @@ -568,19 +505,19 @@ def _download_bulk_parameters( base_url = dataset_config["base_url"] dataset_name = data_type_config["dataset"] - try: - # Determine output file path: base_path/product/dataset/buoy_id/filename.csv - buoy_dir = os.path.join( - self.base_path_to_download, self.product, dataset_name, buoy_id - ) - output_file = os.path.join(buoy_dir, f"buoy_{buoy_id}_bulk_parameters.csv") - - # Check if file exists - if self._check_file_exists(output_file, result, force, dry_run): - return self.finalize_download_result(result) + # Determine output file path: base_path/product/dataset/buoy_id/filename.csv + output_file = os.path.join( + self.base_path_to_download, + self.product, + dataset_name, + buoy_id, + f"buoy_{buoy_id}_bulk_parameters.csv", + ) + missing_years = [] - # Prepare download tasks - download_tasks = [] + def fetch(part_file: str) -> None: + missing_years.clear() + all_data = [] for year in years: urls = [ f"{base_url}/{data_type_config['url_pattern'].format(buoy_id=buoy_id, year=year)}" @@ -589,68 +526,32 @@ def _download_bulk_parameters( urls.append( f"{base_url}/{fallback.format(buoy_id=buoy_id, year=year)}" ) + df = self._download_single_year_bulk(urls, data_type_config["columns"]) + if df is None: + self.logger.warning( + f"Buoy {buoy_id}: No data available for year {year}" + ) + missing_years.append(year) + else: + self.logger.info(f"Buoy {buoy_id}: Data found for year {year}") + all_data.append(df) - download_tasks.append( - { - "urls": urls, - "columns": data_type_config["columns"], - "year": year, - "buoy_id": buoy_id, - } - ) + if not all_data: + raise RemoteFileNotFound(f"No data found for buoy {buoy_id}") - if dry_run: - # In dry run mode, just mark what would be downloaded - for task in download_tasks: - result.add_skipped( - output_file, - f"Would download year {task['year']} (dry run)", - ) - return self.finalize_download_result(result) + combined_df = pd.concat(all_data, ignore_index=True) + combined_df = combined_df.sort_values(["YYYY", "MM", "DD", "hh"]) + combined_df.to_csv(part_file, index=False) - # Execute downloads sequentially - all_data = [] - for task in download_tasks: - try: - df = self._download_single_year_bulk(task["urls"], task["columns"]) - if df is not None: - all_data.append(df) - self.logger.info( - f"Buoy {buoy_id}: Data found for year {task['year']}" - ) - else: - self.logger.warning( - f"Buoy {buoy_id}: No data available for year {task['year']}" - ) - result.add_error( - output_file, - Exception(f"No data available for year {task['year']}"), - ) - except Exception as e: - self.logger.error(f"Error downloading year {task['year']}: {e}") - result.add_error(output_file, e) - - if all_data: - # Combine all years - combined_df = pd.concat(all_data, ignore_index=True) - combined_df = combined_df.sort_values(["YYYY", "MM", "DD", "hh"]) - - # Save to CSV - os.makedirs(buoy_dir, exist_ok=True) - combined_df.to_csv(output_file, index=False) - self.logger.info(f"Data saved to {output_file}") - result.add_downloaded(output_file) - else: - self.logger.error(f"No data found for buoy {buoy_id}") - result.add_error( - output_file, - Exception(f"No data found for buoy {buoy_id}"), - ) - except Exception as e: - result.add_error(output_file, e) - self.logger.error(f"Error processing data for buoy {buoy_id}: {e}") + self.fetch_file( + output_file, fetch, result, force=force, dry_run=dry_run, validate=False + ) + if missing_years: + result.metadata.setdefault("missing_years", {})[output_file] = list( + missing_years + ) - return self.finalize_download_result(result) + return result def _download_single_year_bulk( self, @@ -679,8 +580,7 @@ def _download_single_year_bulk( for url in urls: try: # Download the file - response = requests.get(url, timeout=30) - response.raise_for_status() + response = http_get(self._session, url) content = gzip.decompress(response.content).decode("utf-8") # Skip the header rows and read the data @@ -787,53 +687,25 @@ def _download_wave_spectra( "wave_spectra", ) - if not dry_run: - os.makedirs(buoy_dir, exist_ok=True) - for year in years: url = f"{base_url}/{data_type_config['url_pattern'].format(buoy_id=buoy_id, year=year)}" output_file = os.path.join(buoy_dir, f"buoy_{buoy_id}_spectra_{year}.csv") - # Check if file exists - if self._check_file_exists(output_file, result, force, dry_run): - continue - - if dry_run: - result.add_skipped(output_file, f"Would download year {year} (dry run)") - continue - - try: - # Download and read the data - response = requests.get(url, timeout=30) - response.raise_for_status() - - # Read the data + def fetch(part_file: str, url: str = url, year: int = year) -> None: + response = http_get(self._session, url) df = pd.read_csv( io.BytesIO(response.content), compression="gzip", sep=r"\s+", na_values=["MM", "99.00", "999.0"], ) - - # Skip if empty or invalid data if df.empty or len(df.columns) < 5: - self.logger.warning(f"No valid data for {buoy_id} - {year}") - result.add_error( - output_file, - Exception(f"No valid data for {buoy_id} - {year}"), - context={"year": year}, - ) - continue - - # Save the data - df.to_csv(output_file, index=False) - result.add_downloaded(output_file) - self.logger.info(f"Successfully saved data for {buoy_id} - {year}") + raise RemoteFileNotFound(f"No valid data for {buoy_id} - {year}") + df.to_csv(part_file, index=False) - except Exception as e: - self.logger.warning(f"No data found for: {buoy_id} - {year}: {e}") - result.add_error(output_file, e, context={"year": year}) - continue + self.fetch_file( + output_file, fetch, result, force=force, dry_run=dry_run, validate=False + ) return result @@ -890,8 +762,6 @@ def _download_directional_spectra( buoy_id, "directional_spectra", ) - if not dry_run: - os.makedirs(buoy_dir, exist_ok=True) for year in years: for coef, info in coefficients.items(): @@ -899,127 +769,9 @@ def _download_directional_spectra( url = f"{base_url}/{info['url_pattern'].format(buoy_id=buoy_id, year=year)}" save_path = os.path.join(buoy_dir, filename) - # Check if file exists - if self._check_file_exists(save_path, result, force, dry_run): - continue - - if dry_run: - result.add_skipped( - save_path, - f"Would download {info['name']} for year {year} (dry run)", - ) - continue - - try: - self.logger.debug(f"Downloading {info['name']} data for {year}...") - - # Download the file - response = requests.get(url, stream=True, timeout=30) - response.raise_for_status() - - # Save the compressed file - with open(save_path, "wb") as f: - shutil.copyfileobj(response.raw, f) - - result.add_downloaded(save_path) - self.logger.info(f"Successfully downloaded {filename}") - - except Exception as e: - self.logger.warning(f"Error downloading {filename}: {e}") - result.add_error(save_path, e) - continue - - return self.finalize_download_result(result) - - def _download_wind_forecast( - self, - data_type_config: dict, - dataset_config: dict, - date: str = None, - region: List[float] = None, - force: bool = False, - dry_run: bool = False, - ) -> DownloadResult: - """ - Download NOAA GFS wind forecast data. - - Downloads and crops GFS wind forecast data for a specific date and region. - Files are saved to: - base_path_to_download/product/dataset/{date}_{region}.nc - - Parameters - ---------- - data_type_config : dict - Configuration for the data type. - dataset_config : dict - Configuration for the dataset. - date : str, optional - Date to download data for (format: "YYYYMMDD"). If None, uses today's date. - Default is None. - region : List[float], optional - Geographic region coordinates. Default is None. - force : bool, optional - Force re-download even if file exists. Default is False. - dry_run : bool, optional - If True, only check what would be downloaded. Default is False. - - Returns - ------- - DownloadResult - Result with information about downloaded, skipped, and error files. + def fetch(part_file: str, url: str = url) -> None: + http_download(self._session, url, part_file, resume=False) - Notes - ----- - This method will be DEPRECATED in the future. - """ + self.fetch_file(save_path, fetch, result, force=force, dry_run=dry_run) - if date is None: - date = datetime.today().strftime("%Y%m%d") - - self.logger.info(f"Downloading wind forecast for date {date}") - - result = self.create_download_result() - url_base = dataset_config["base_url"] - dataset_name = data_type_config["dataset"] - dbn = "gfs_0p25_1hr" - url = f"{url_base}/gfs{date}/{dbn}_00z" - - # File path for local storage: base_path/product/dataset/filename.nc - forecast_dir = os.path.join( - self.base_path_to_download, self.product, dataset_name - ) - if not dry_run: - os.makedirs(forecast_dir, exist_ok=True) - - file_path = os.path.join( - forecast_dir, f"{date}_{'_'.join(map(str, region))}.nc" - ) - - # Check if file exists - if self._check_file_exists(file_path, result, force, dry_run): - return result - - if dry_run: - result.add_skipped( - file_path, f"Would download wind forecast for {date} (dry run)" - ) - return result - - try: - self.logger.info(f"Downloading and cropping forecast data from: {url}") - # Open dataset from URL - data = xr.open_dataset(url) - - # Select only wind data - variables = data_type_config["variables"] - data_select = data[variables] - - self.logger.info(f"Storing local copy at: {file_path}") - data_select.to_netcdf(file_path) - result.add_downloaded(file_path) - - except Exception as e: - self.logger.error(f"Error downloading wind forecast: {e}") - result.add_error(file_path, e) - - return self.finalize_download_result(result) + return result diff --git a/docs/source/downloaders.rst b/docs/source/downloaders.rst new file mode 100644 index 0000000..612999e --- /dev/null +++ b/docs/source/downloaders.rst @@ -0,0 +1,280 @@ +Data downloaders +================ + +``bluemath_tk.downloaders`` downloads climate and ocean data from several +providers with a common interface: every downloader can do a dry run first, +skips files that are already complete, retries failed transfers, and reports +what happened in a :class:`~bluemath_tk.downloaders.DownloadResult`. + +Providers and products +---------------------- + +.. list-table:: + :header-rows: 1 + :widths: 15 25 15 45 + + * - Provider + - Class + - Products + - Notes + * - ``csiro`` + - ``CSIRODownloader`` + - WHACS, CAWCR + - WAVEWATCH III hindcasts from the CSIRO THREDDS server. Full monthly files + (resumable, parallel) or time series at given points (OPeNDAP). + * - ``copernicus`` + - ``CopernicusDownloader`` + - ERA5, CERRA + - Copernicus Climate Data Store (``cdsapi``). Parallel requests. + * - ``noaa`` + - ``NOAADownloader`` + - NDBC + - NDBC buoy bulk parameters, wave spectra and directional spectra. + * - ``aviso`` + - ``AvisoDownloader`` + - SWOT + - AVISO+ FTP server (requires an account). + * - ``ecmwf`` + - ``ECMWFDownloader`` + - OpenData + - ECMWF IFS/AIFS real-time forecasts (``ecmwf-opendata``), from the ECMWF + portal or its AWS, Azure and Google mirrors. + +The Copernicus and ECMWF downloaders need the optional dependencies: + +.. code-block:: bash + + pip install bluemath-tk[downloaders] + +Package layout +-------------- + +Each provider is a subpackage with a single downloader module and one folder per +product, which only holds the product configuration (servers, datasets, +variables, file names): + +.. code-block:: text + + bluemath_tk/downloaders/ + ├── _base_downloaders.py # BaseDownloader: routing, fetch_file, results + ├── _download_result.py # DownloadResult + ├── _transfer.py # retries, resumable HTTP, file validation + └── csiro/ + ├── __init__.py # exports CSIRODownloader + ├── csiro_downloader.py + ├── CAWCR/CAWCR_config.json + └── WHACS/WHACS_config.json + +Downloaders are imported from their provider subpackage: + +.. code-block:: python + + from bluemath_tk.downloaders.csiro import CSIRODownloader + from bluemath_tk.downloaders.copernicus import CopernicusDownloader + +Common behaviour +---------------- + +All downloaders share these rules: + +``dry_run`` + Defaults to ``True``: nothing is downloaded and the result lists what would + be. Pass ``dry_run=False`` to download. + +Complete files only + Each file is written to ``.part``, checked (NetCDF files must open, + GRIB files must end with the GRIB end marker, HTTP transfers must match the + remote size) and then renamed. A file without ``.part`` is always complete, + so an interrupted run never leaves a file that looks finished. + +Re-runs are cheap + Existing files are skipped (``force=True`` downloads them again). CSIRO files + are also compared with the remote size (``verify_existing=True``), and + interrupted CSIRO transfers resume where they stopped. + +Retries + Each file is tried ``retries`` times (default 3), waiting ``retry_backoff`` + seconds (default 30, doubled each time). Files that are not published yet + are reported as skipped, not as errors, so the same call can be repeated to + keep a dataset up to date. + +Results + Every call returns a :class:`~bluemath_tk.downloaders.DownloadResult`: + + .. code-block:: python + + print(result.summary()) # counts, skip reasons, size to download, errors + result.downloaded_files # list of paths + result.has_errors + result.save("run.json") # keep a log of the run + +CSIRO: WHACS and CAWCR +---------------------- + +Full monthly files, one per variable for WHACS: + +.. code-block:: python + + from bluemath_tk.downloaders.csiro import CSIRODownloader + + downloader = CSIRODownloader( + product="WHACS", + base_path_to_download="/lustre/geocean/DATA/CSIRO", + max_workers=4, # parallel downloads + ) + downloader.list_datasets() # ['global-smc'] + downloader.list_variables("global-smc") + + result = downloader.download_data( + dataset="global-smc", + variables=["phs0", "ptp0", "pdir0"], + start="2020-01", + end="2023-12", + dry_run=True, + ) + print(result.summary()) # includes the total size to download + +CAWCR has two datasets, ``gridded`` (all variables in one file per month) and +``spec`` (spectra at output stations). Their station list and variable names +changed in June 2013; the downloader stores both periods in separate folders +and renames the old variables to the new ones (``Efth`` → ``efth``, ``depth`` → +``dpt``, ...). + +Time series at one or more points are extracted through OPeNDAP, without +downloading the full files. One file is written per month and, by default, all +months are combined into one file: + +.. code-block:: python + + downloader = CSIRODownloader(product="CAWCR", base_path_to_download="./data") + result = downloader.download_point( + dataset="spec", + lon=-4.0, + lat=44.0, + start="2013-01", + end="2013-12", + dry_run=False, + ) + result.metadata["combined_file"] + +The nearest model point is used; its distance is stored in the file attributes +(``nearest_point_distance_km``) and a warning is logged above 25 km. OPeNDAP +extraction is slow for long periods (the server reads whole files), so for many +years it is faster to download the files and read them locally. + +Copernicus: ERA5 and CERRA +-------------------------- + +The CDS key is read from ``~/.cdsapirc`` (or the ``CDSAPI_KEY`` environment +variable) unless ``api_key`` is given. + +.. code-block:: python + + from bluemath_tk.downloaders.copernicus import CopernicusDownloader + + downloader = CopernicusDownloader( + product="CERRA", + base_path_to_download="/lustre/geocean/DATA/Copernicus", + max_workers=4, # CDS requests in parallel + ) + result = downloader.download_data( + variables=["2m_temperature", "mean_sea_level_pressure"], + years=range(2020, 2026), + months=range(1, 13), + file_period="month", + dry_run=True, + ) + +Files are saved to +``//////``. The +``file_period`` argument decides how requests are split: + +- ``"month"``: one file per variable and month, ``__.nc``. +- ``"year"``: one file per variable and year, ``_.nc``. +- ``"auto"`` (default): yearly files when all 12 months are requested, + monthly files otherwise. + +Yearly files with the previous naming (``__01_02_..._12.nc``) are +recognised and not downloaded again. + +NOAA, AVISO and ECMWF +--------------------- + +.. code-block:: python + + from bluemath_tk.downloaders.noaa import NOAADownloader + + downloader = NOAADownloader(product="NDBC", base_path_to_download="./data") + downloader.list_data_types() + result = downloader.download_data( + data_type="wave_spectra", buoy_id="41001", years=[2022, 2023], dry_run=False + ) + +.. code-block:: python + + from bluemath_tk.downloaders.aviso import AvisoDownloader + + with AvisoDownloader( + product="SWOT", + base_path_to_download="./swot", + username="user", + password="password", + ) as downloader: + result = downloader.download_data( + dataset="swot-l3-expert", cycles=["cycle_001"], dry_run=False + ) + +.. code-block:: python + + from bluemath_tk.downloaders.ecmwf import ECMWFDownloader + + downloader = ECMWFDownloader( + product="OpenData", + base_path_to_download="./ecmwf", + source="aws", # mirror; "ecmwf" is rate limited + ) + result = downloader.download_data( + dataset="forecast_data", param=["msl"], step=[0, 24], type="fc", dry_run=False + ) + +ECMWF files are saved per forecast run, +``/OpenData/forecast_data////``. Without +``date`` the latest run is used; ``date=-1, time=12`` selects yesterday's 12 UTC +run. + +Adding a product or a provider +------------------------------ + +A new product of an existing provider only needs a configuration file: + +1. Create ``//_config.json``, following the files of + the other products of that provider. +2. Add it to the ``products_configs`` dictionary of the downloader class. +3. If the product needs its own logic, add a ``download_data_`` + method (snake case: ``OpenData`` → ``download_data_open_data``); + ``download_data`` routes to it. + +A new provider is a new subpackage with the same layout. Its class inherits +from :class:`~bluemath_tk.downloaders.BaseDownloader`, declares +``products_configs``, and writes every file through ``fetch_file``, which +handles skipping, dry runs, ``.part`` files, validation and retries: + +.. code-block:: python + + def download_data_myproduct(self, ..., force=False, dry_run=True): + result = self.create_download_result() + for url, output_file in files: + def fetch(part_file, url=url): + http_download(self.session, url, part_file) + self.fetch_file(output_file, fetch, result, force=force, dry_run=dry_run) + return self.finalize_download_result(result) + +The ``tests/downloaders/test_structure.py`` test checks that every provider +follows this layout and that ``products_configs`` matches the product folders. + +Automated downloads at GeoOcean +------------------------------- + +On the GeoOcean cluster, download jobs for ``/lustre/geocean/DATA`` live in +``/lustre/geocean/DATA/DownloadMe``: one small script and Slurm job per dataset, +which only says *what* to download, plus their logs. See its ``README.md``. diff --git a/docs/source/index.rst b/docs/source/index.rst index 2fe10f2..4de8b87 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -14,6 +14,7 @@ BlueMath-tk: A Python Library for Coastal Climate Hazards installation contribute modules + downloaders validation benchmarking diff --git a/tests/downloaders/test_copernicus.py b/tests/downloaders/test_copernicus.py new file mode 100644 index 0000000..f3999c5 --- /dev/null +++ b/tests/downloaders/test_copernicus.py @@ -0,0 +1,114 @@ +import os +import tempfile +import unittest + +try: + from bluemath_tk.downloaders.copernicus import CopernicusDownloader +except ImportError: # cdsapi is an optional dependency + CopernicusDownloader = None + + +@unittest.skipIf(CopernicusDownloader is None, "cdsapi not installed") +class TestCopernicusTasks(unittest.TestCase): + """ + Requests and file names, checked with dry runs (no request is sent to the CDS). + """ + + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + self.cerra = CopernicusDownloader( + product="CERRA", + base_path_to_download=self.tmp_dir.name, + api_key="dummy", + debug=False, + ) + self.era5 = CopernicusDownloader( + product="ERA5", + base_path_to_download=self.tmp_dir.name, + api_key="dummy", + debug=False, + ) + + def tearDown(self): + self.tmp_dir.cleanup() + + def _tasks(self, downloader, **kwargs): + tasks = [] + downloader._run_tasks = lambda new_tasks, force, dry_run: tasks.extend( + new_tasks + ) + downloader.download_data(**kwargs) + return tasks + + @staticmethod + def _names(tasks): + return [os.path.basename(task["output_file"]) for task in tasks] + + def test_file_periods(self): + full_year = dict( + variables=["2m_temperature"], years=[2020], months=range(1, 13) + ) + self.assertEqual( + self._names(self._tasks(self.cerra, **full_year)), + ["2m_temperature_2020.nc"], + ) + monthly = self._names(self._tasks(self.cerra, file_period="month", **full_year)) + self.assertEqual(len(monthly), 12) + self.assertEqual(monthly[0], "2m_temperature_2020_01.nc") + self.assertEqual( + self._names( + self._tasks( + self.cerra, variables=["2m_temperature"], years=[2020], months=[3] + ) + ), + ["2m_temperature_2020_03.nc"], + ) + with self.assertRaises(ValueError): + self._tasks(self.cerra, file_period="week", **full_year) + + def test_cerra_request(self): + (task,) = self._tasks( + self.cerra, variables=["mean_sea_level_pressure"], years=[2024], months=[2] + ) + self.assertEqual(task["dataset"], "reanalysis-cerra-single-levels") + self.assertEqual(task["request"]["month"], ["02"]) + self.assertEqual(task["request"]["year"], ["2024"]) + self.assertEqual(len(task["request"]["time"]), 8) + + def test_era5_spectra_dates(self): + (task,) = self._tasks( + self.era5, + variables=["spectra"], + years=[2024], + months=[2], + area=[45, -5, 43, -3], + ) + self.assertEqual(task["request"]["date"], "2024-02-01/to/2024-02-29") + self.assertEqual(task["request"]["area"], "45/-5/43/-3") + + def test_legacy_yearly_file_is_recognised(self): + (task,) = self._tasks( + self.cerra, variables=["2m_temperature"], years=[2020], months=range(1, 13) + ) + legacy_file = task["legacy_file"] + os.makedirs(os.path.dirname(legacy_file)) + open(legacy_file, "w").close() + + downloader = CopernicusDownloader( + product="CERRA", + base_path_to_download=self.tmp_dir.name, + api_key="dummy", + debug=False, + ) + result = downloader.download_data( + variables=["2m_temperature"], + years=[2020], + months=list(range(1, 13)), + file_period="month", + dry_run=True, + ) + self.assertEqual(result.skipped_files, [legacy_file]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/downloaders/test_csiro.py b/tests/downloaders/test_csiro.py new file mode 100644 index 0000000..7fc247d --- /dev/null +++ b/tests/downloaders/test_csiro.py @@ -0,0 +1,148 @@ +import os +import unittest + +import numpy as np +import pandas as pd +from test_transfer import LocalServerTestCase + +from bluemath_tk.downloaders.csiro.csiro_downloader import ( + CSIRODownloader, + nearest_point_index, +) + + +class TestNearestPointIndex(unittest.TestCase): + def test_longitude_conventions(self): + lons = np.array([0.0, 90.0, 355.0, 180.0]) + lats = np.array([0.0, 0.0, 44.0, 0.0]) + idx, distance = nearest_point_index(lons, lats, -5.0, 44.0) + self.assertEqual(idx, 2) + self.assertAlmostEqual(distance, 0.0, places=3) + + def test_distance(self): + idx, distance = nearest_point_index(np.array([0.0]), np.array([0.0]), 1.0, 0.0) + self.assertEqual(idx, 0) + self.assertAlmostEqual(distance, 111.19, delta=0.1) + + +class TestCSIROPaths(unittest.TestCase): + def setUp(self): + self.whacs = CSIRODownloader(product="WHACS", base_path_to_download="/data") + self.cawcr = CSIRODownloader(product="CAWCR", base_path_to_download="/data") + + def test_unknown_product(self): + with self.assertRaises(ValueError): + CSIRODownloader(product="OTHER", base_path_to_download="/data") + + def test_whacs_paths(self): + period = pd.Period("2000-02", freq="M") + url = self.whacs.get_url("global-smc", period, "phs1") + self.assertTrue( + url.startswith("https://data-cbr.csiro.au/thredds5/fileServer/catch_all/") + ) + self.assertTrue( + url.endswith( + "/global-smc/1hr/phs1/" + "phs1_WHACS_hindcast_WHACS_ERA5_1hr_200002010000-200002292300.nc" + ) + ) + self.assertEqual( + self.whacs.get_local_path("global-smc", period, "phs1"), + os.path.normpath( + "/data/WHACS/global-smc/phs1/" + "phs1_WHACS_hindcast_WHACS_ERA5_1hr_200002010000-200002292300.nc" + ), + ) + self.assertIn( + "/thredds5/dodsC/", + self.whacs.get_url("global-smc", period, "phs1", service="opendap"), + ) + + def test_cawcr_spec_eras(self): + self.assertEqual( + self.cawcr.get_local_path("spec", pd.Period("2013-05", freq="M")), + os.path.normpath("/data/CAWCR/CSIRO_spectra/1979/ww3.201305_spec.nc"), + ) + self.assertEqual( + self.cawcr.get_local_path("spec", pd.Period("2013-06", freq="M")), + os.path.normpath("/data/CAWCR/CSIRO_spectra/2013/ww3.201306_spec.nc"), + ) + self.assertEqual( + self.cawcr.get_local_path("gridded", pd.Period("2025-01", freq="M")), + os.path.normpath("/data/CAWCR/CSIRO_partition/ww3.glob_24m.202501.nc"), + ) + + def test_months_are_clipped_to_availability(self): + months = self.whacs.get_months("global-smc", "1970-01", "2030-12") + self.assertEqual(str(months[0]), "1979-01") + self.assertEqual(str(months[-1]), "2023-12") + self.assertEqual( + len(self.whacs.get_months("global-smc", "2000-01", "2000-12")), 12 + ) + + def test_variables_are_checked(self): + with self.assertRaises(ValueError): + self.whacs.download_data("global-smc") + with self.assertRaises(ValueError): + self.whacs.download_data("global-smc", variables=["nope"]) + with self.assertRaises(ValueError): + self.whacs.download_data("unknown-dataset", variables=["hs"]) + + +class TestCSIRODownloadFile(LocalServerTestCase): + def setUp(self): + super().setUp() + self.downloader = CSIRODownloader( + product="CAWCR", + base_path_to_download=self.tmp_dir.name, + retries=1, + retry_backoff=0, + ) + self.url = f"{self.base_url}/data.bin" + self.output_file = os.path.join(self.tmp_dir.name, "data.bin") + + def test_download_and_skip(self): + result = self.downloader._download_file( + self.url, self.output_file, force=False, verify_existing=True, dry_run=False + ) + self.assertEqual(result.downloaded_files, [self.output_file]) + with open(self.output_file, "rb") as f: + self.assertEqual(f.read(), self.content) + + result = self.downloader._download_file( + self.url, self.output_file, force=False, verify_existing=True, dry_run=False + ) + self.assertEqual(result.skipped_files, [self.output_file]) + + def test_incomplete_local_file_is_downloaded_again(self): + with open(self.output_file, "wb") as f: + f.write(self.content[:100]) + + result = self.downloader._download_file( + self.url, self.output_file, force=False, verify_existing=True, dry_run=True + ) + self.assertEqual(result.skipped_files, [self.output_file]) + self.assertIn( + "would download again", result.metadata["skip_reasons"][self.output_file] + ) + + result = self.downloader._download_file( + self.url, self.output_file, force=False, verify_existing=True, dry_run=False + ) + self.assertEqual(result.downloaded_files, [self.output_file]) + self.assertEqual(os.path.getsize(self.output_file), len(self.content)) + + def test_not_available_on_server(self): + result = self.downloader._download_file( + f"{self.base_url}/missing.nc", + os.path.join(self.tmp_dir.name, "missing.nc"), + force=False, + verify_existing=True, + dry_run=False, + ) + self.assertEqual(len(result.skipped_files), 1) + self.assertFalse(result.has_errors) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/downloaders/test_noaa.py b/tests/downloaders/test_noaa.py new file mode 100644 index 0000000..c63e2dc --- /dev/null +++ b/tests/downloaders/test_noaa.py @@ -0,0 +1,32 @@ +import os +import tempfile +import unittest + +import pandas as pd + +from bluemath_tk.downloaders.noaa.noaa_downloader import read_bulk_parameters + + +class TestReadBulkParameters(unittest.TestCase): + def test_years_are_not_duplicated(self): + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp_dir: + buoy_dir = os.path.join(tmp_dir, "NDBC", "buoy_data", "41001") + os.makedirs(buoy_dir) + pd.DataFrame( + { + "YYYY": [2020, 2021, 2022], + "MM": [1, 1, 1], + "DD": [1, 1, 1], + "hh": [0, 0, 0], + "mm": [0, 0, 0], + "WVHT": [1.0, 2.0, 3.0], + } + ).to_csv(os.path.join(buoy_dir, "buoy_41001_bulk_parameters.csv")) + + df = read_bulk_parameters(tmp_dir, "41001", [2020, 2021]) + self.assertEqual(df["WVHT"].tolist(), [1.0, 2.0]) + self.assertIsNone(read_bulk_parameters(tmp_dir, "41001", 1999)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/downloaders/test_structure.py b/tests/downloaders/test_structure.py new file mode 100644 index 0000000..6fc1280 --- /dev/null +++ b/tests/downloaders/test_structure.py @@ -0,0 +1,88 @@ +import importlib +import json +import os +import unittest + +import bluemath_tk.downloaders as downloaders + +DOWNLOADERS_DIR = os.path.dirname(downloaders.__file__) + +# provider package -> downloader class +PROVIDERS = { + "aviso": "AvisoDownloader", + "copernicus": "CopernicusDownloader", + "csiro": "CSIRODownloader", + "ecmwf": "ECMWFDownloader", + "noaa": "NOAADownloader", +} + + +def _files(path): + return sorted( + name + for name in os.listdir(path) + if name != "__pycache__" and not name.startswith(".") + ) + + +class TestDownloadersLayout(unittest.TestCase): + """ + Each provider folder holds only __init__.py, _downloader.py and one + folder per product with _config.json. + """ + + def test_providers_are_known(self): + providers = [ + name + for name in _files(DOWNLOADERS_DIR) + if os.path.isdir(os.path.join(DOWNLOADERS_DIR, name)) + ] + self.assertEqual(sorted(providers), sorted(PROVIDERS)) + + def test_provider_folders(self): + for provider in PROVIDERS: + with self.subTest(provider=provider): + path = os.path.join(DOWNLOADERS_DIR, provider) + files = [ + name + for name in _files(path) + if os.path.isfile(os.path.join(path, name)) + ] + self.assertEqual( + files, sorted(["__init__.py", f"{provider}_downloader.py"]) + ) + for product in self._product_dirs(provider): + self.assertEqual( + _files(os.path.join(path, product)), + [f"{product}_config.json"], + ) + with open( + os.path.join(path, product, f"{product}_config.json") + ) as f: + self.assertIsInstance(json.load(f), dict) + + def test_products_match_config_folders(self): + for provider, class_name in PROVIDERS.items(): + with self.subTest(provider=provider): + try: + module = importlib.import_module( + f"bluemath_tk.downloaders.{provider}" + ) + except ImportError as e: + self.skipTest(f"optional dependency missing: {e}") + downloader_class = getattr(module, class_name) + self.assertEqual( + sorted(downloader_class.list_products()), + self._product_dirs(provider), + ) + + @staticmethod + def _product_dirs(provider): + path = os.path.join(DOWNLOADERS_DIR, provider) + return [ + name for name in _files(path) if os.path.isdir(os.path.join(path, name)) + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/downloaders/test_transfer.py b/tests/downloaders/test_transfer.py new file mode 100644 index 0000000..5209595 --- /dev/null +++ b/tests/downloaders/test_transfer.py @@ -0,0 +1,331 @@ +import os +import tempfile +import threading +import unittest +from functools import partial +from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer + +import numpy as np +import xarray as xr + +from bluemath_tk.downloaders import BaseDownloader, DownloadResult +from bluemath_tk.downloaders._transfer import ( + RemoteFileNotFound, + http_download, + make_session, + remote_size, + retry_call, + validate_file, +) + + +class RangeRequestHandler(SimpleHTTPRequestHandler): + """Static file handler with support for "Range: bytes=N-" requests.""" + + def log_message(self, *args): + pass + + def do_GET(self): + range_header = self.headers.get("Range") + path = self.translate_path(self.path) + if not range_header or not os.path.isfile(path): + return super().do_GET() + start = int(range_header.split("=")[1].split("-")[0]) + with open(path, "rb") as f: + data = f.read()[start:] + self.send_response(206) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + +class LocalServerTestCase(unittest.TestCase): + """Serve a temporary directory over HTTP during the tests.""" + + @classmethod + def setUpClass(cls): + cls.served_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + handler = partial(RangeRequestHandler, directory=cls.served_dir.name) + cls.server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + threading.Thread(target=cls.server.serve_forever, daemon=True).start() + cls.base_url = f"http://127.0.0.1:{cls.server.server_address[1]}" + cls.content = os.urandom(3 * 1024 * 1024 + 123) + with open(os.path.join(cls.served_dir.name, "data.bin"), "wb") as f: + f.write(cls.content) + open(os.path.join(cls.served_dir.name, "empty.bin"), "wb").close() + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + cls.server.server_close() + cls.served_dir.cleanup() + + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + self.session = make_session(retries=0) + + def tearDown(self): + self.tmp_dir.cleanup() + + +class TestHttpDownload(LocalServerTestCase): + def test_full_download(self): + target = os.path.join(self.tmp_dir.name, "data.bin") + url = f"{self.base_url}/data.bin" + http_download(self.session, url, target, expected_size=len(self.content)) + with open(target, "rb") as f: + self.assertEqual(f.read(), self.content) + + def test_resume_partial_file(self): + target = os.path.join(self.tmp_dir.name, "data.bin") + with open(target, "wb") as f: + f.write(self.content[:1000]) + http_download( + self.session, + f"{self.base_url}/data.bin", + target, + expected_size=len(self.content), + ) + with open(target, "rb") as f: + self.assertEqual(f.read(), self.content) + + def test_size_mismatch_raises(self): + target = os.path.join(self.tmp_dir.name, "data.bin") + with self.assertRaises(IOError): + http_download( + self.session, f"{self.base_url}/data.bin", target, expected_size=10 + ) + + def test_missing_file(self): + url = f"{self.base_url}/missing.bin" + with self.assertRaises(RemoteFileNotFound): + remote_size(self.session, url) + with self.assertRaises(RemoteFileNotFound): + http_download(self.session, url, os.path.join(self.tmp_dir.name, "x")) + + def test_empty_remote_file_is_missing(self): + with self.assertRaises(RemoteFileNotFound): + remote_size(self.session, f"{self.base_url}/empty.bin") + + def test_remote_size(self): + size = remote_size(self.session, f"{self.base_url}/data.bin") + self.assertEqual(size, len(self.content)) + + +class TestRetryCall(unittest.TestCase): + def test_succeeds_after_failures(self): + calls = [] + + def flaky(): + calls.append(1) + if len(calls) < 3: + raise ConnectionError("temporary") + return "ok" + + self.assertEqual(retry_call(flaky, retries=3, backoff=0), "ok") + self.assertEqual(len(calls), 3) + + def test_gives_up(self): + def failing(): + raise ConnectionError("permanent") + + with self.assertRaises(ConnectionError): + retry_call(failing, retries=2, backoff=0) + + def test_no_retry_for_missing_files(self): + calls = [] + + def missing(): + calls.append(1) + raise RemoteFileNotFound("missing") + + with self.assertRaises(RemoteFileNotFound): + retry_call(missing, retries=3, backoff=0) + self.assertEqual(len(calls), 1) + + +class TestValidateFile(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + + def tearDown(self): + self.tmp_dir.cleanup() + + def _write_netcdf(self) -> str: + path = os.path.join(self.tmp_dir.name, "data.nc") + xr.Dataset({"hs": ("time", np.arange(1000.0))}).to_netcdf( + path, engine="netcdf4" + ) + return path + + def test_valid_netcdf(self): + validate_file(self._write_netcdf()) + + def test_truncated_netcdf(self): + path = self._write_netcdf() + with open(path, "rb") as f: + data = f.read() + with open(path, "wb") as f: + f.write(data[: len(data) // 2]) + with self.assertRaises(IOError): + validate_file(path) + + def test_empty_file(self): + path = os.path.join(self.tmp_dir.name, "empty.nc") + open(path, "wb").close() + with self.assertRaises(IOError): + validate_file(path) + + def test_grib_end_marker(self): + path = os.path.join(self.tmp_dir.name, "data.grib") + with open(path, "wb") as f: + f.write(b"GRIB" + b"\x00" * 100 + b"7777") + validate_file(path) + with open(path, "wb") as f: + f.write(b"GRIB" + b"\x00" * 100) + with self.assertRaises(IOError): + validate_file(path) + + +class DummyDownloader(BaseDownloader): + products_configs = { + "MyProduct": {"datasets": {"a": {}, "b": {}}}, + "Other": {"datasets": {}}, + } + + def __init__(self, base_path_to_download: str, product: str = "MyProduct"): + super().__init__( + product=product, + base_path_to_download=base_path_to_download, + retries=2, + retry_backoff=0, + ) + + def download_data_my_product(self, value, dry_run=True): + return value, dry_run + + +class TestFetchFile(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + self.downloader = DummyDownloader(self.tmp_dir.name) + self.output_file = os.path.join(self.tmp_dir.name, "sub", "file.txt") + self.result = DownloadResult() + + def tearDown(self): + self.tmp_dir.cleanup() + + @staticmethod + def _write(part_file: str) -> None: + with open(part_file, "w") as f: + f.write("data") + + def test_download(self): + self.downloader.fetch_file( + self.output_file, self._write, self.result, dry_run=False + ) + self.assertEqual(self.result.downloaded_files, [self.output_file]) + self.assertTrue(os.path.exists(self.output_file)) + self.assertFalse(os.path.exists(self.output_file + ".part")) + + def test_existing_file_is_skipped(self): + os.makedirs(os.path.dirname(self.output_file)) + self._write(self.output_file) + self.downloader.fetch_file( + self.output_file, self._write, self.result, dry_run=False + ) + self.assertEqual(self.result.skipped_files, [self.output_file]) + self.assertEqual(self.result.downloaded_files, []) + + def test_dry_run_writes_nothing(self): + self.downloader.fetch_file(self.output_file, self._write, self.result) + self.assertEqual(self.result.skipped_files, [self.output_file]) + self.assertFalse(os.path.exists(os.path.dirname(self.output_file))) + + def test_failure_leaves_no_file(self): + calls = [] + + def failing(part_file): + calls.append(1) + with open(part_file, "w") as f: + f.write("half") + raise ConnectionError("interrupted") + + self.downloader.fetch_file( + self.output_file, failing, self.result, dry_run=False + ) + self.assertEqual(len(calls), 2) + self.assertEqual(self.result.error_files, [self.output_file]) + self.assertFalse(os.path.exists(self.output_file)) + self.assertFalse(os.path.exists(self.output_file + ".part")) + + def test_resumable_failure_keeps_part_file(self): + def failing(part_file): + with open(part_file, "a") as f: + f.write("half") + raise ConnectionError("interrupted") + + self.downloader.fetch_file( + self.output_file, failing, self.result, dry_run=False, resumable=True + ) + self.assertFalse(os.path.exists(self.output_file)) + self.assertTrue(os.path.exists(self.output_file + ".part")) + + def test_invalid_file_is_not_kept(self): + def empty(part_file): + open(part_file, "w").close() + + self.downloader.fetch_file(self.output_file, empty, self.result, dry_run=False) + self.assertEqual(self.result.error_files, [self.output_file]) + self.assertFalse(os.path.exists(self.output_file)) + self.assertFalse(os.path.exists(self.output_file + ".part")) + + def test_missing_remote_file_is_skipped(self): + def missing(part_file): + raise RemoteFileNotFound("http://example.com/file.txt") + + self.downloader.fetch_file( + self.output_file, missing, self.result, dry_run=False + ) + self.assertEqual(self.result.skipped_files, [self.output_file]) + self.assertEqual(self.result.error_files, []) + + +class TestBaseDownloader(unittest.TestCase): + def test_routing_to_product_method(self): + downloader = DummyDownloader(tempfile.gettempdir()) + self.assertEqual(downloader.download_data(1), (1, True)) + self.assertEqual(downloader.download_data(2, dry_run=False), (2, False)) + self.assertEqual(downloader.list_datasets(), ["a", "b"]) + + def test_unknown_product(self): + with self.assertRaises(ValueError): + DummyDownloader(tempfile.gettempdir(), product="Missing") + self.assertEqual(DummyDownloader.list_products(), ["MyProduct", "Other"]) + + def test_product_without_download_method(self): + downloader = DummyDownloader(tempfile.gettempdir(), product="Other") + with self.assertRaises(ValueError): + downloader.download_data() + + +class TestDownloadResult(unittest.TestCase): + def test_merge_and_save(self): + result = DownloadResult(downloaded_files=["a"]) + other = DownloadResult(downloaded_files=["a", "b"]) + other.add_skipped("c", "Already downloaded") + other.add_error("d", ValueError("bad")) + result.merge(other) + self.assertEqual(result.downloaded_files, ["a", "b"]) + self.assertEqual(result.skipped_files, ["c"]) + self.assertEqual(result.error_files, ["d"]) + self.assertEqual(result.metadata["skip_reasons"], {"c": "Already downloaded"}) + + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp_dir: + path = os.path.join(tmp_dir, "logs", "result.json") + result.save(path) + self.assertTrue(os.path.exists(path)) + + +if __name__ == "__main__": + unittest.main()