diff --git a/.github/workflows/build-documentation.yml b/.github/workflows/build-documentation.yml index 45cbc8b..ebc7223 100644 --- a/.github/workflows/build-documentation.yml +++ b/.github/workflows/build-documentation.yml @@ -13,10 +13,10 @@ jobs: steps: - uses: actions/checkout@v3 - - name: Set up Python 3.10 + - name: Set up Python 3.11 uses: actions/setup-python@v4 with: - python-version: '3.10' + python-version: '3.11' - name: Install dependencies run: | sudo apt-get update diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml index 3e7c6b1..e0c9bf1 100644 --- a/.github/workflows/linting.yml +++ b/.github/workflows/linting.yml @@ -15,7 +15,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.10', '3.11', '3.12'] + python-version: ['3.11', '3.12', '3.13'] steps: - uses: actions/checkout@v3 diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml index e0b69d8..bc1f8f9 100644 --- a/.github/workflows/publish-to-pypi.yml +++ b/.github/workflows/publish-to-pypi.yml @@ -29,7 +29,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v4 with: - python-version: '3.10' + python-version: '3.11' - name: Install dependencies run: | python -m pip install --upgrade pip diff --git a/.github/workflows/smoke-test.yml b/.github/workflows/smoke-test.yml index 1182210..bd0ef0b 100644 --- a/.github/workflows/smoke-test.yml +++ b/.github/workflows/smoke-test.yml @@ -13,7 +13,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.10', '3.11', '3.12'] + python-version: ['3.11', '3.12', '3.13'] steps: - uses: actions/checkout@v3 diff --git a/.github/workflows/testing-and-coverage.yml b/.github/workflows/testing-and-coverage.yml index 7f845dc..e2fea00 100644 --- a/.github/workflows/testing-and-coverage.yml +++ b/.github/workflows/testing-and-coverage.yml @@ -15,7 +15,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.10', '3.11', '3.12'] + python-version: ['3.11', '3.12', '3.13'] steps: - uses: actions/checkout@v3 diff --git a/docs/notebooks/pzserver_tutorial.ipynb b/docs/notebooks/pzserver_tutorial.ipynb index 967b338..a8bf94a 100644 --- a/docs/notebooks/pzserver_tutorial.ipynb +++ b/docs/notebooks/pzserver_tutorial.ipynb @@ -59,9 +59,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "scrolled": true - }, + "metadata": {}, "outputs": [], "source": [ "# pip install --upgrade --no-cache-dir pzserver" @@ -200,7 +198,6 @@ "cell_type": "code", "execution_count": null, "metadata": { - "scrolled": true, "tags": [] }, "outputs": [], @@ -218,9 +215,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "scrolled": true - }, + "metadata": {}, "outputs": [], "source": [ "pz_server.display_products_list(filters={\"release\": \"DP0.2\", \n", @@ -237,9 +232,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "scrolled": true - }, + "metadata": {}, "outputs": [], "source": [ "pz_server.display_products_list(filters={\"product_type\": \"estimates\"})" @@ -256,7 +249,6 @@ "cell_type": "code", "execution_count": null, "metadata": { - "scrolled": true, "tags": [] }, "outputs": [], @@ -765,9 +757,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "scrolled": true - }, + "metadata": {}, "outputs": [], "source": [ "# Show all products\n", @@ -1016,9 +1006,7 @@ { "cell_type": "code", "execution_count": null, - "metadata": { - "scrolled": true - }, + "metadata": {}, "outputs": [], "source": [ "# Show all products\n", diff --git a/pyproject.toml b/pyproject.toml index 3f88d40..078e69c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,17 +6,13 @@ readme = "README.md" authors = [ { name = "LIneA", email = "itteam@linea.org.br" } ] -requires-python = ">=3.7,<3.14" +requires-python = ">=3.11,<3.14.9" classifiers = [ "Development Status :: 4 - Beta", "License :: OSI Approved :: MIT License", "Intended Audience :: Developers", "Intended Audience :: Science/Research", "Operating System :: OS Independent", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", diff --git a/requirements.txt b/requirements.txt index ff8e2ea..82c9b69 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,10 +3,11 @@ Deprecated>=1,<2 h5py>=3,<4 ipython>=8,<9 Jinja2>=3,<4 +lsdb>=0.9,<1 matplotlib>=3,<4 numpy>=1,<3 pandas>=2,<3 pyarrow>=17 pydantic>=2,<3 requests>=2,<3 -tables-io>=1,<2 \ No newline at end of file +tables-io>=1,<2 diff --git a/src/.pylintrc b/src/.pylintrc index 81771e9..8ab5949 100644 --- a/src/.pylintrc +++ b/src/.pylintrc @@ -83,7 +83,7 @@ limit-inference-results=100 load-plugins= # Pickle collected data for later comparisons. -persistent=yes +persistent=no # Minimum Python version to use for version dependent checks. Will default to # the version used to run pylint. @@ -92,9 +92,7 @@ py-version=3.9 # Discover python modules and packages in the file system subtree. recursive=no -# When enabled, pylint would attempt to guess common misconfiguration and emit -# user-friendly hints instead of false-positive error messages. -suggestion-mode=yes +# suggestion-mode was removed from newer pylint versions. # Allow loading of arbitrary C extensions. Extensions are imported into the # active Python interpreter and may run arbitrary code. diff --git a/src/pzserver/communicate.py b/src/pzserver/communicate.py index 2a85b61..52cf8cd 100644 --- a/src/pzserver/communicate.py +++ b/src/pzserver/communicate.py @@ -3,11 +3,16 @@ """ import json +import pathlib +import re +import time +from email.message import Message +from urllib.parse import urljoin import requests -# pylint: disable=R0904 +# pylint: disable=too-many-lines,too-many-public-methods class PzRequests: """ Responsible for managing all requests to the Pz Server app. @@ -170,6 +175,7 @@ def _check_response(self, api_response) -> dict: def _send_request( self, prerequest, + *, stream=False, timeout=None, verify=True, @@ -320,40 +326,116 @@ def _check_token(self): msg = cntxt.get("message", "Unforeseen error") raise requests.exceptions.RequestException(f"Status code {stcode}: {msg}") - def _download_request(self, url, save_in="."): + @staticmethod + def _filename_from_content_disposition(content_disposition): + message = Message() + message["content-disposition"] = content_disposition or "" + filename = message.get_filename() + if filename: + return filename + + return "download" + + @staticmethod + def _download_total_size(response, fallback_size=None): + content_range = response.headers.get("Content-Range", "") + match = re.match(r"bytes \d+-\d+/(\d+)$", content_range) + if match: + return int(match.group(1)) + + content_length = response.headers.get("Content-Length") + if content_length is None: + return fallback_size + + content_length = int(content_length) + if response.status_code == 206 and fallback_size is not None: + return fallback_size + content_length + return content_length + + def _download_response(self, url, start_byte=None): + headers = {"Authorization": f"Token {self._token}"} + if start_byte: + headers["Range"] = f"bytes={start_byte}-" + + req = requests.Request("GET", url, headers=headers) + return self._send_request(req.prepare(), stream=True) + + def _download_request( # pylint: disable=too-many-locals + self, url, save_in=".", max_attempts=3 + ): """ Download a record from the API. Args: url (str): url to get save_in (str): location where the file will be saved + max_attempts (int): number of attempts when streaming is interrupted """ - req = requests.Request( - "GET", - url, - headers=dict( - { - "Authorization": f"Token {self._token}", - } - ), - ) - data = self._send_request(req.prepare(), stream=True) + save_dir = pathlib.Path(save_in) + save_dir.mkdir(parents=True, exist_ok=True) + filename = None + destination = None + partial_destination = None + last_error = None + + for _attempt in range(max_attempts): + start_byte = ( + partial_destination.stat().st_size + if partial_destination and partial_destination.exists() + else None + ) + data = self._download_response(url, start_byte=start_byte) + if not data.get("success", False): + return data - if data.get("success", False): resp_obj = data.get("response_object", None) - filename = resp_obj.headers.get("Content-Disposition", "") - filename = filename.split("filename=")[1] - filename = f"{save_in}/{filename}" + if filename is None: + content_disposition = resp_obj.headers.get("Content-Disposition", "") + filename = self._filename_from_content_disposition(content_disposition) + destination = save_dir / filename + partial_destination = destination.with_name(f"{destination.name}.part") + + if partial_destination.exists(): + resp_obj.close() + continue + + if resp_obj.status_code != 206 and partial_destination.exists(): + partial_destination.unlink() + start_byte = None + + mode = "ab" if resp_obj.status_code == 206 and start_byte else "wb" + expected_size = self._download_total_size(resp_obj, start_byte) + + try: + with open(partial_destination, mode) as filedown: + for chunk in resp_obj.iter_content(chunk_size=1024 * 1024): + if chunk: + filedown.write(chunk) + except requests.exceptions.RequestException as error: + last_error = error + continue + finally: + resp_obj.close() + + actual_size = partial_destination.stat().st_size + if expected_size is None or actual_size == expected_size: + partial_destination.replace(destination) + data.update({"message": str(destination)}) + return data + + last_error = requests.exceptions.ChunkedEncodingError( + f"Incomplete download: {actual_size} of {expected_size} bytes" + ) - with open(filename, "wb") as filedown: - for chunk in resp_obj.iter_content(chunk_size=128): - filedown.write(chunk) + partial_path = str(partial_destination) if partial_destination else "unknown" + raise requests.exceptions.RequestException( + f"Download interrupted after {max_attempts} attempts. " + f"Partial file kept at: {partial_path}" + ) from last_error - data.update({"message": filename}) - return data - - return data + def _resolve_api_url(self, url): + return urljoin(self._base_api_url, url) def _patch_request(self, url, data) -> dict: """ @@ -603,20 +685,33 @@ def options(self, entity) -> dict: return opt.get("data") - def download_main_file(self, _id, save_in="."): + def download_main_file(self, _id, save_in=".", timeout=1800, poll_interval=2): """ Gets the contents uploaded by the user for a given record. Args: _id (int): record id save_in (str): location where the file will be saved + timeout (int): maximum seconds to wait for archive preparation + poll_interval (int): seconds between status checks Returns: dict: record data """ + archive = self._prepare_product_main_file_download(_id) + archive = self._wait_for_product_download_ready( + _id, + archive, + timeout=timeout, + poll_interval=poll_interval, + status_getter=self._get_product_main_file_download_status, + download_name="Product main file", + ) + return self._download_request( - f"{self._base_api_url}products/{_id}/download_main_file/", save_in + self._resolve_api_url(archive["download_url"]), + save_in, ) def get_product_files(self, product_id) -> list: @@ -704,20 +799,113 @@ def get_main_file_info(self, _id) -> dict: return data - def download_product(self, _id, save_in="."): + def _prepare_product_download(self, _id): + data = self._post_request( + f"{self._base_api_url}products/{_id}/download/prepare/", + payload=None, + ) + + if "success" in data and data["success"] is False: + raise requests.exceptions.RequestException(data["message"]) + + return data.get("data") + + def _get_product_download_status(self, _id): + data = self._get_request( + f"{self._base_api_url}products/{_id}/download/status/" + ) + + if "success" in data and data["success"] is False: + raise requests.exceptions.RequestException(data["message"]) + + return data.get("data") + + def _prepare_product_main_file_download(self, _id): + data = self._post_request( + f"{self._base_api_url}products/{_id}/download/main-file/prepare/", + payload=None, + ) + + if "success" in data and data["success"] is False: + raise requests.exceptions.RequestException(data["message"]) + + return data.get("data") + + def _get_product_main_file_download_status(self, _id): + data = self._get_request( + f"{self._base_api_url}products/{_id}/download/main-file/status/" + ) + + if "success" in data and data["success"] is False: + raise requests.exceptions.RequestException(data["message"]) + + return data.get("data") + + def _wait_for_product_download_ready( + self, + _id, + archive, + timeout=1800, + poll_interval=2, + *, + status_getter=None, + download_name="Product download", + ): + deadline = time.monotonic() + timeout + status_getter = status_getter or self._get_product_download_status + + while archive and archive.get("status") in ("pending", "running"): + if time.monotonic() >= deadline: + raise requests.exceptions.RequestException( + f"{download_name} is still being prepared. Please try again later." + ) + + time.sleep(poll_interval) + archive = status_getter(_id) + + if archive and archive.get("status") == "failed": + raise requests.exceptions.RequestException( + archive.get("error_message") or f"Failed to prepare {download_name}." + ) + + if not archive or archive.get("status") != "ready": + raise requests.exceptions.RequestException( + f"{download_name} archive is not ready." + ) + + download_url = archive.get("download_url") + if not download_url: + raise requests.exceptions.RequestException( + f"{download_name} archive is ready but no download URL was returned." + ) + + return archive + + def download_product(self, _id, save_in=".", timeout=1800, poll_interval=2): """ Downloads the product to local Args: _id (int): record id save_in (str): location where the file will be saved + timeout (int): maximum seconds to wait for archive preparation + poll_interval (int): seconds between status checks Returns: dict: record data """ + archive = self._prepare_product_download(_id) + archive = self._wait_for_product_download_ready( + _id, + archive, + timeout=timeout, + poll_interval=poll_interval, + ) + return self._download_request( - f"{self._base_api_url}products/{_id}/download/", save_in + self._resolve_api_url(archive["download_url"]), + save_in, ) def start_process(self, data, files=None): diff --git a/src/pzserver/core.py b/src/pzserver/core.py index d98d371..5d67f89 100644 --- a/src/pzserver/core.py +++ b/src/pzserver/core.py @@ -2,10 +2,14 @@ Classes responsible for managing user interaction """ +import importlib +import pathlib import tempfile import time +import zipfile import pandas as pd +import requests import tables_io from astropy.table import Table from IPython.display import display @@ -22,6 +26,11 @@ FONTCOLORERR = "\033[38;2;255;0;0m" FONTCOLOREND = "\033[0m" +HATS_DIRECTORY_MESSAGE = ( + "This product main file is a directory-based dataset, such as a " + "HATS collection." +) +MAX_IN_MEMORY_PRODUCT_SIZE_KB = 200 * 1024 class PzServer: @@ -350,7 +359,7 @@ def download_product(self, product_id=None, save_in="."): else: print(f"{FONTCOLORERR}Error: {results_dict['message']}{FONTCOLORERR}") - def get_product(self, product_id=None): + def get_product(self, product_id=None, get_big_products=False): """ Fetches the data product contents to local. @@ -362,6 +371,9 @@ def get_product(self, product_id=None): product_id (str or int): data product unique identifier (product id number or internal name) + get_big_products (bool, optional): allow products whose + stored main file is larger than 200 MB to be loaded + into memory. Defaults to False. Returns: astropy.Table @@ -369,6 +381,7 @@ def get_product(self, product_id=None): """ print("Connecting to PZ Server...") metadata = self.get_product_metadata(product_id) + self._validate_product_size(metadata, product_id, get_big_products) prod_type = metadata["product_type_internal_name"] if prod_type in ("validation_results", "training_results"): @@ -382,16 +395,25 @@ def get_product(self, product_id=None): if not metadata["main_file"]: raise FileNotFoundError(f"Product ID ({product_id}): main file not found") + if metadata["main_file"].get("is_directory"): + return self._get_hats_product(metadata) + file_extension = metadata["main_file"]["extension"] with tempfile.TemporaryDirectory() as tmpdirname: results_dict = self.api.download_main_file(metadata["id"], tmpdirname) if not results_dict.get("success", False): - print(f"Error: {results_dict['message']}") - return None + message = results_dict.get("message", "Failed to download main file.") + raise requests.exceptions.RequestException(message) + + file_path = pathlib.Path(results_dict["message"]) + if file_path.suffix.lower() == ".zip": + return self._read_hats_archive( + file_path, + metadata["main_file"].get("name"), + ) - file_path = results_dict["message"] if file_extension == ".csv": # TBD: add CSV to tables_io supported formats delimiter = metadata["main_file"].get("delimiter", None) @@ -413,7 +435,159 @@ def get_product(self, product_id=None): table = tables_io.read(file_path, tables_io.types.AP_TABLE) return table - def upload( + @staticmethod + def _validate_product_size(metadata, product_id, get_big_products): + """Prevent large products from being loaded into memory by default.""" + main_file = metadata.get("main_file") or {} + size_kb = main_file.get("size") + try: + is_too_large = float(size_kb) > MAX_IN_MEMORY_PRODUCT_SIZE_KB + except (TypeError, ValueError): + is_too_large = False + + if not is_too_large or get_big_products: + return + + product_name = metadata.get("internal_name") or product_id + raise ValueError( + f"Product '{product_name}' is larger than 200 MB and cannot be loaded " + "into memory by default. PZ Server may store products in compressed " + "form, so this product can use substantially more RAM when decompressed. " + "We recommend downloading it and reading it locally with pandas, Dask, " + "or LSDB (for HATS catalogs):\n\n" + "from pathlib import Path\n\n" + f"prod_name = {product_name!r}\n" + 'caminho = Path("./downloaded_data")\n' + "caminho.mkdir(parents=True, exist_ok=True)\n" + "pz_server.download_product(product_id=prod_name, save_in=caminho)\n\n" + "To continue anyway, call " + "get_product(product_id=prod_name, get_big_products=True)." + ) + + def _get_hats_product(self, metadata): + """Download a directory-based HATS main file and load it with LSDB.""" + lsdb = self._import_lsdb() + + with tempfile.TemporaryDirectory() as tmpdirname: + archive_path = self._download_hats_main_file(metadata, tmpdirname) + return self._read_hats_archive( + archive_path, + metadata["main_file"].get("name"), + lsdb=lsdb, + ) + + def _download_hats_main_file(self, metadata, destination): + results_dict = self.api.download_main_file(metadata["id"], destination) + if results_dict.get("success", False): + return pathlib.Path(results_dict["message"]) + + raise requests.exceptions.RequestException( + results_dict.get("message", "Failed to download HATS main file.") + ) + + @staticmethod + def _import_lsdb(): + try: + return importlib.import_module("lsdb") + except ImportError as exc: + raise ImportError( + f"{HATS_DIRECTORY_MESSAGE} Reading it with get_product() requires " + "the optional dependency 'lsdb'. Install it and try again." + ) from exc + + def _read_hats_archive(self, archive_path, main_file_name, lsdb=None): + if lsdb is None: + lsdb = self._import_lsdb() + + with tempfile.TemporaryDirectory() as tmpdirname: + extracted_dir = pathlib.Path(tmpdirname, "extracted") + self._extract_zip_safely(archive_path, extracted_dir) + + catalog = self._open_hats_catalog_from_directory( + lsdb, + extracted_dir, + main_file_name, + ) + return self._lsdb_catalog_to_table(catalog) + + @staticmethod + def _extract_zip_safely(archive_path, destination): + """Extract a zip archive while preventing paths outside destination.""" + destination = pathlib.Path(destination).resolve() + destination.mkdir(parents=True, exist_ok=True) + + with zipfile.ZipFile(archive_path) as archive: + for member in archive.infolist(): + target = (destination / member.filename).resolve() + try: + target.relative_to(destination) + except ValueError as exc: + raise ValueError( + f"Unsafe path in product archive: {member.filename}" + ) from exc + + archive.extractall(destination) + + def _open_hats_catalog_from_directory(self, lsdb, extracted_dir, main_file_name): + errors = [] + for candidate in self._hats_open_candidates(extracted_dir, main_file_name): + try: + return lsdb.open_catalog(path=str(candidate)) + except Exception as exc: # pylint: disable=broad-exception-caught + errors.append(f"{candidate}: {exc}") + + raise ValueError( + f"{HATS_DIRECTORY_MESSAGE} Could not open it with LSDB. Tried: " + + "; ".join(errors) + ) + + def _hats_open_candidates(self, extracted_dir, main_file_name): + extracted_dir = pathlib.Path(extracted_dir) + candidates = [] + + if main_file_name: + candidates.append(extracted_dir / main_file_name) + + for filename in ("collection.properties", "hats.properties", "properties"): + candidates.extend(path.parent for path in extracted_dir.rglob(filename)) + + candidates.append(extracted_dir) + + unique_candidates = [] + seen = set() + for candidate in candidates: + candidate = candidate.resolve() + if candidate in seen or not candidate.is_dir(): + continue + seen.add(candidate) + unique_candidates.append(candidate) + + return unique_candidates + + @staticmethod + def _lsdb_catalog_to_table(catalog): + if hasattr(catalog, "compute"): + data = catalog.compute() + elif hasattr(catalog, "to_pandas"): + data = catalog.to_pandas() + else: + ddf = getattr(catalog, "_ddf", None) + if ddf is None: + raise TypeError( + "Could not convert LSDB catalog to an astropy Table." + ) + data = ddf.compute() + + if isinstance(data, Table): + return data + if isinstance(data, pd.DataFrame): + return Table.from_pandas(data) + if hasattr(data, "to_pandas"): + return Table.from_pandas(data.to_pandas()) + + raise TypeError("Could not convert LSDB catalog data to an astropy Table.") + + def upload( # pylint: disable=too-many-positional-arguments self, name: str, product_type: str, diff --git a/tests/.pylintrc b/tests/.pylintrc index c3cf2bf..3f80ffc 100644 --- a/tests/.pylintrc +++ b/tests/.pylintrc @@ -83,7 +83,7 @@ limit-inference-results=100 load-plugins= # Pickle collected data for later comparisons. -persistent=yes +persistent=no # Minimum Python version to use for version dependent checks. Will default to # the version used to run pylint. @@ -92,9 +92,7 @@ py-version=3.9 # Discover python modules and packages in the file system subtree. recursive=no -# When enabled, pylint would attempt to guess common misconfiguration and emit -# user-friendly hints instead of false-positive error messages. -suggestion-mode=yes +# suggestion-mode was removed from newer pylint versions. # Allow loading of arbitrary C extensions. Extensions are imported into the # active Python interpreter and may run arbitrary code. diff --git a/tests/pzserver/test_download.py b/tests/pzserver/test_download.py new file mode 100644 index 0000000..b00ec10 --- /dev/null +++ b/tests/pzserver/test_download.py @@ -0,0 +1,239 @@ +""" +Tests for product download orchestration. +""" + +import importlib.util +from pathlib import Path + +import pytest +import requests + + +def load_communicate_module(): + module_path = Path(__file__).parents[2] / "src" / "pzserver" / "communicate.py" + spec = importlib.util.spec_from_file_location("pzserver_communicate", module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def make_api(): + communicate = load_communicate_module() + api = object.__new__(communicate.PzRequests) + api._base_api_url = "https://pz.example.org/api/" + api._token = "token" + return api, communicate + + +class FakeDownloadResponse: + """Minimal streaming response fake used by download tests.""" + + def __init__(self, headers, chunks, status_code=200, error=None): + self.headers = headers + self._chunks = chunks + self.status_code = status_code + self.error = error + self.closed = False + + def iter_content(self, chunk_size): # pylint: disable=unused-argument + yield from self._chunks + if self.error: + raise self.error + + def close(self): + self.closed = True + + +def test_download_product_uses_prepared_download_url(monkeypatch): + api, _ = make_api() + calls = [] + + monkeypatch.setattr( + api, + "_prepare_product_download", + lambda product_id: { + "status": "ready", + "download_url": "/api/products/42/download/file/?token=signed", + }, + ) + monkeypatch.setattr( + api, + "_download_request", + lambda url, save_in: calls.append((url, save_in)) + or {"success": True, "message": f"{save_in}/product.zip"}, + ) + + result = api.download_product(42, "/tmp") + + assert result["success"] is True + assert calls == [ + ("https://pz.example.org/api/products/42/download/file/?token=signed", "/tmp") + ] + + +def test_download_product_polls_until_ready(monkeypatch): + api, communicate = make_api() + statuses = [ + {"status": "running"}, + { + "status": "ready", + "download_url": "/api/products/42/download/file/?token=signed", + }, + ] + + monkeypatch.setattr( + api, + "_prepare_product_download", + lambda product_id: {"status": "pending"}, + ) + monkeypatch.setattr( + api, + "_get_product_download_status", + lambda product_id: statuses.pop(0), + ) + monkeypatch.setattr( + api, + "_download_request", + lambda url, save_in: {"success": True}, + ) + monkeypatch.setattr(communicate.time, "sleep", lambda seconds: None) + + result = api.download_product(42, "/tmp", timeout=10, poll_interval=0) + + assert result["success"] is True + assert not statuses + + +def test_download_product_raises_when_archive_failed(monkeypatch): + api, _ = make_api() + + monkeypatch.setattr( + api, + "_prepare_product_download", + lambda product_id: {"status": "failed", "error_message": "zip failed"}, + ) + + with pytest.raises(requests.exceptions.RequestException, match="zip failed"): + api.download_product(42) + + +def test_download_main_file_uses_prepared_download_url(monkeypatch): + api, _ = make_api() + calls = [] + + monkeypatch.setattr( + api, + "_prepare_product_main_file_download", + lambda product_id: { + "status": "ready", + "download_url": "/api/products/42/download/main-file/file/?token=signed", + }, + ) + monkeypatch.setattr( + api, + "_download_request", + lambda url, save_in: calls.append((url, save_in)) + or {"success": True, "message": f"{save_in}/main_file.csv"}, + ) + + result = api.download_main_file(42, "/tmp") + + assert result["success"] is True + assert calls == [ + ( + "https://pz.example.org/api/products/42/download/main-file/file/?token=signed", + "/tmp", + ) + ] + + +def test_download_main_file_polls_until_ready(monkeypatch): + api, communicate = make_api() + statuses = [ + {"status": "running"}, + { + "status": "ready", + "download_url": "/api/products/42/download/main-file/file/?token=signed", + }, + ] + + monkeypatch.setattr( + api, + "_prepare_product_main_file_download", + lambda product_id: {"status": "pending"}, + ) + monkeypatch.setattr( + api, + "_get_product_main_file_download_status", + lambda product_id: statuses.pop(0), + ) + monkeypatch.setattr( + api, + "_download_request", + lambda url, save_in: {"success": True}, + ) + monkeypatch.setattr(communicate.time, "sleep", lambda seconds: None) + + result = api.download_main_file(42, "/tmp", timeout=10, poll_interval=0) + + assert result["success"] is True + assert not statuses + + +def test_download_main_file_raises_when_archive_failed(monkeypatch): + api, _ = make_api() + + monkeypatch.setattr( + api, + "_prepare_product_main_file_download", + lambda product_id: {"status": "failed", "error_message": "zip failed"}, + ) + + with pytest.raises(requests.exceptions.RequestException, match="zip failed"): + api.download_main_file(42) + + +def test_download_request_resumes_interrupted_stream(tmp_path): + api, _ = make_api() + calls = [] + first_response = FakeDownloadResponse( + { + "Content-Disposition": "attachment; filename=product.zip", + "Content-Length": "6", + }, + [b"abc"], + error=requests.exceptions.ChunkedEncodingError("connection broken"), + ) + second_response = FakeDownloadResponse( + { + "Content-Disposition": "attachment; filename=product.zip", + "Content-Range": "bytes 3-5/6", + "Content-Length": "3", + }, + [b"def"], + status_code=206, + ) + responses = [first_response, second_response] + + def send_request(prepared, stream=False): + calls.append((prepared.headers.copy(), stream)) + return { + "success": True, + "response_object": responses.pop(0), + } + + api._send_request = send_request + + result = api._download_request("https://pz.example.org/file", tmp_path) + + assert result["success"] is True + assert Path(result["message"]).read_bytes() == b"abcdef" + assert not Path(f"{result['message']}.part").exists() + assert calls[0] == ({"Authorization": "Token token"}, True) + assert calls[1] == ( + {"Authorization": "Token token", "Range": "bytes=3-"}, + True, + ) + assert first_response.closed is True + assert second_response.closed is True + assert not responses diff --git a/tests/pzserver/test_get_product.py b/tests/pzserver/test_get_product.py new file mode 100644 index 0000000..8319f93 --- /dev/null +++ b/tests/pzserver/test_get_product.py @@ -0,0 +1,268 @@ +""" +Tests for product retrieval behavior. +""" + +import importlib.util +import sys +import zipfile +from pathlib import Path +from types import ModuleType +from unittest import mock + +import pandas as pd +import pytest +import requests +from astropy.table import Table + + +def load_core_module(): + root = Path(__file__).parents[2] / "src" / "pzserver" + package = ModuleType("pzserver") + package.__path__ = [str(root)] + + sys.modules.setdefault("pzserver", package) + spec = importlib.util.spec_from_file_location( + "pzserver.core", + root / "core.py", + submodule_search_locations=[str(root)], + ) + module = importlib.util.module_from_spec(spec) + sys.modules["pzserver.core"] = module + spec.loader.exec_module(module) + return module + + +def make_server_with_api(core, api): + server = object.__new__(core.PzServer) + server.api = api + return server + + +def base_metadata(main_file): + return { + "id": 42, + "product_type_internal_name": "redshift_catalog", + "main_file": main_file, + } + + +def test_get_product_rejects_products_larger_than_200_mb(): + core = load_core_module() + api = mock.Mock() + server = make_server_with_api(core, api) + metadata = base_metadata( + { + "extension": ".parquet", + "is_directory": False, + "name": "main.parquet", + "size": 200 * 1024 + 1, + } + ) + metadata["internal_name"] = "874_test_crc" + server.get_product_metadata = mock.Mock(return_value=metadata) + + with pytest.raises(ValueError, match="larger than 200 MB") as exc_info: + server.get_product("874_test_crc") + + message = str(exc_info.value) + assert "pz_server.download_product" in message + assert "pandas, Dask, or LSDB" in message + assert "get_big_products=True" in message + api.download_main_file.assert_not_called() + + +def test_get_product_allows_large_products_when_explicitly_requested(tmp_path): + core = load_core_module() + table_path = tmp_path / "large.parquet" + api = mock.Mock() + api.download_main_file.return_value = { + "success": True, + "message": str(table_path), + } + server = make_server_with_api(core, api) + server.get_product_metadata = mock.Mock( + return_value=base_metadata( + { + "extension": ".parquet", + "is_directory": False, + "name": "large.parquet", + "size": 200 * 1024 + 1, + } + ) + ) + + expected = Table({"id": [1]}) + with mock.patch.object(core.tables_io, "read", return_value=expected): + result = server.get_product("large", get_big_products=True) + + assert result is expected + api.download_main_file.assert_called_once() + + +def test_get_product_allows_product_at_exactly_200_mb(tmp_path): + core = load_core_module() + table_path = tmp_path / "limit.parquet" + api = mock.Mock() + api.download_main_file.return_value = { + "success": True, + "message": str(table_path), + } + server = make_server_with_api(core, api) + server.get_product_metadata = mock.Mock( + return_value=base_metadata( + { + "extension": ".parquet", + "is_directory": False, + "name": "limit.parquet", + "size": 200 * 1024, + } + ) + ) + + expected = Table({"id": [1]}) + with mock.patch.object(core.tables_io, "read", return_value=expected): + result = server.get_product("limit") + + assert result is expected + api.download_main_file.assert_called_once() + + +class FakeLsdbCatalog: # pylint: disable=too-few-public-methods + """Minimal fake LSDB catalog used by get_product tests.""" + def compute(self): + return pd.DataFrame( + { + "object_id": [1, 2], + "ra": [10.0, 20.0], + "dec": [-10.0, -20.0], + } + ) + + +def create_hats_archive(path, root=""): + prefix = f"{root.rstrip('/')}" if root else "" + prefix = f"{prefix}/" if prefix else "" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr(f"{prefix}collection.properties", "catalog_name=test\n") + archive.writestr(f"{prefix}dataset/Norder=0/Dir=0/Npix=0.parquet", b"fake") + + +def test_get_product_downloads_hats_main_file_and_reads_with_lsdb(tmp_path, monkeypatch): + core = load_core_module() + fake_lsdb = ModuleType("lsdb") + open_catalog = mock.Mock(return_value=FakeLsdbCatalog()) + fake_lsdb.open_catalog = open_catalog + monkeypatch.setitem(sys.modules, "lsdb", fake_lsdb) + + archive_path = tmp_path / "hats.zip" + create_hats_archive(archive_path) + + api = mock.Mock() + api.download_main_file.return_value = { + "success": True, + "message": str(archive_path), + } + server = make_server_with_api(core, api) + server.get_product_metadata = mock.Mock( + return_value=base_metadata( + { + "extension": "", + "is_directory": True, + "name": "main", + } + ) + ) + + table = server.get_product("hats_product") + + assert isinstance(table, Table) + assert table.colnames == ["object_id", "ra", "dec"] + assert len(table) == 2 + api.download_main_file.assert_called_once() + api.download_product.assert_not_called() + assert open_catalog.call_args.kwargs["path"].endswith("/extracted") + + +def test_get_product_reads_zip_download_even_when_metadata_is_not_directory( + tmp_path, monkeypatch +): + core = load_core_module() + fake_lsdb = ModuleType("lsdb") + open_catalog = mock.Mock(return_value=FakeLsdbCatalog()) + fake_lsdb.open_catalog = open_catalog + monkeypatch.setitem(sys.modules, "lsdb", fake_lsdb) + + archive_path = tmp_path / "hats.zip" + create_hats_archive(archive_path) + + api = mock.Mock() + api.download_main_file.return_value = { + "success": True, + "message": str(archive_path), + } + server = make_server_with_api(core, api) + server.get_product_metadata = mock.Mock( + return_value=base_metadata( + { + "extension": ".parquet", + "is_directory": False, + "name": "main", + } + ) + ) + + table = server.get_product("hats_product") + + assert isinstance(table, Table) + assert table.colnames == ["object_id", "ra", "dec"] + api.download_main_file.assert_called_once() + api.download_product.assert_not_called() + assert open_catalog.call_args.kwargs["path"].endswith("/extracted") + + +def test_get_product_raises_when_hats_main_file_download_fails(monkeypatch): + core = load_core_module() + fake_lsdb = ModuleType("lsdb") + fake_lsdb.open_catalog = mock.Mock(return_value=FakeLsdbCatalog()) + monkeypatch.setitem(sys.modules, "lsdb", fake_lsdb) + + api = mock.Mock() + api.download_main_file.return_value = { + "success": False, + "message": "HATS main file download failed", + } + server = make_server_with_api(core, api) + server.get_product_metadata = mock.Mock( + return_value=base_metadata( + { + "extension": "", + "is_directory": True, + "name": "main", + } + ) + ) + + with pytest.raises(requests.exceptions.RequestException, match="HATS"): + server.get_product("hats_product") + + api.download_main_file.assert_called_once() + api.download_product.assert_not_called() + + +def test_get_product_requires_lsdb_for_hats(monkeypatch): + core = load_core_module() + monkeypatch.setitem(sys.modules, "lsdb", None) + api = mock.Mock() + server = make_server_with_api(core, api) + server.get_product_metadata = mock.Mock( + return_value=base_metadata( + { + "extension": "", + "is_directory": True, + "name": "main", + } + ) + ) + + with pytest.raises(ImportError, match="lsdb"): + server.get_product("hats_product")