diff --git a/config/simdb.cfg b/config/simdb.cfg index 4a6816f5..5315fab0 100644 --- a/config/simdb.cfg +++ b/config/simdb.cfg @@ -34,3 +34,4 @@ result_backend = redis://redis:6379/0 [partition] data = /data/simdb/partition +http = /data/simdb/http \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index eff71ff3..f9b57500 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -20,6 +20,7 @@ services: - ./validation:/app/validation:ro - ./config:/app/config:ro - ./tmp/partition_data:/data/simdb/partition:ro + - ./tmp/http:/data/simdb/http - ./upload_folder:/data/simdb/simulations depends_on: redis: @@ -38,6 +39,7 @@ services: volumes: - ./config:/app/config:ro - ./tmp/partition_data:/data/simdb/partition:ro + - ./tmp/http:/data/simdb/http - ./upload_folder:/data/simdb/simulations depends_on: redis: diff --git a/docs/generate_cli_docs.py b/docs/generate_cli_docs.py index faf5527e..b0a00933 100644 --- a/docs/generate_cli_docs.py +++ b/docs/generate_cli_docs.py @@ -39,7 +39,9 @@ def extract_sub_commands(output: str) -> list[str]: sub_commands = [] for line in output.split("\n"): if in_commands: - if line: + # Lines indented further than the command column are continuations + # of the previous command's help text. + if line and not line.startswith(" "): sub_commands.append(line.split()[0]) if line == "Commands:": in_commands = True diff --git a/docs/how-to/push-pull.md b/docs/how-to/push-pull.md index 6b555080..2b021989 100644 --- a/docs/how-to/push-pull.md +++ b/docs/how-to/push-pull.md @@ -42,6 +42,108 @@ simdb simulation push SIM_ID --add-watcher See [watchers](../explanation/concepts.md#watchers) and the `simdb remote watcher` commands in the [CLI reference](../reference/cli.md). +### Resumable uploads + +Against a v1.3 remote, `push` sends the file bytes using the IETF +[Resumable Uploads for HTTP](https://datatracker.ietf.org/doc/draft-ietf-httpbis-resumable-upload/) +protocol (draft 11, interop version 8). No shared file system is required. The +server stages the uploaded bytes in its `http` partition, and from there the +flow matches `push_local`: SimDB sends the metadata, the server queues the copy +into its upload folder as a background +[Celery task](operate-server/run-celery-workers.md), and the CLI blocks while +reporting the ingestion state. Against a v1.2 remote, `push` falls back to the +earlier non-resumable transfer. + +Because the protocol is resumable, an interrupted upload (a lost connection, +Ctrl-C) does not have to start over. Re-running `push` for the same simulation +asks the server how many bytes it already holds for each file and continues from +that offset. + +The chunk size is governed by the server: it advertises a maximum append size +through the `Upload-Limit` response header and the client sizes its chunks to +stay within that bound. The limit defaults to 8 MiB and can be tuned with +[`server.max_append_size`](../reference/server-configuration.md#server), for +example to fit within a reverse proxy's request body limit. + +Each chunk is integrity-checked with an RFC 9530 `Content-Digest` (SHA-256), so +the server can verify a chunk before appending it. A digest mismatch is rejected +with a `400` response and the offending bytes are not stored, so corruption in +transit cannot be silently committed. + +### Stage uploads on the server + +`push` needs no client-side partition configuration — the paths come straight +from the local simulation. The server, however, must define where uploaded bytes +are staged, through a partition named `http` in its `simdb.cfg` (see +[Server configuration](../reference/server-configuration.md#partition)): + +```ini +[partition] +http = /var/lib/simdb/http-staging +``` + +A file uploaded to `/file/` is written to +`//file/` and referenced by an +`http:////file/` URI. The background ingestion task resolves that +URI against the `http` partition, copies the file into the simulation's upload +folder, and finally removes the staged copy. + +Subfolder structure is handled just as it is for +[`push_local`](#push-on-a-shared-file-system): files keep their relative layout, +so multi-file IMAS datasets (HDF5, ASCII, and MDSplus backends) stay contained +within their own directory and are reconstructed correctly on the server, while +standalone files (such as an IMAS netCDF `.nc`) are not given a spurious +enclosing folder. + +## Push on a shared file system + +If your machine and the server can reach the same physical file paths (as on the +ITER network), sending large datasets over HTTP is slow and redundant. Use +`push_local` instead: + +```bash +simdb simulation push_local SIM_ID +``` + +`push_local` sends only the metadata and the storage paths. The server then + +1. validates the metadata against the active schemas, +2. queues the file copy as a background [Celery task](operate-server/run-celery-workers.md), and +3. completes the ingestion once the copy finishes. + +The command blocks and reports the ingestion state as it changes: + +```text +Waiting for ingestion to complete... QUEUED -> COPYING -> COPIED -> COMPLETED +Successfully pushed simulation UUID +``` + +### Configure partitions + +For `push_local` to resolve files on both sides, client and server must agree on +a set of *partitions*: short logical names mapped to absolute directories. Add a +`[partition]` section to your client configuration (see +[Client configuration](../reference/configuration.md#partition)): + +```ini +[partition] +data = /home/user/my_simdb_data +work = /work/imas/shared +sdcc = / +``` + +Mapping `sdcc` to the system root makes any path under `/sdcc/projects/...` +match, so `/sdcc/projects/my_run` becomes `sdcc:sdcc/projects/my_run`. When +several partitions contain a file the most specific (deepest) path wins, so a +catch-all mapping like this never shadows the others. + +When you run `push_local`, SimDB checks every input and output path against your +partitions. A path inside a partition is rewritten to a partition-relative URI — +`/home/user/my_simdb_data/scenarios/run1.txt` becomes +`data:scenarios/run1.txt`. The server resolves that URI against its own +`[partition]` configuration, so the two sides may mount the same storage at +different absolute paths. + ## Pull Pull copies a simulation's metadata into your local catalogue and downloads its diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index f2bbc65c..2a4aa04e 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -90,6 +90,17 @@ One section per configured remote server. Manage these with | --- | --- | | `file` | Path to the local SQLite catalogue. Defaults to `sim.db` in the user data directory (for example `~/.local/share/simdb/sim.db`). | +### `[partition]` + +Maps logical partition names to absolute directories on this machine. Used by +`simdb simulation push_local` to rewrite file paths into partition-relative URIs +that the server can resolve (see +[Push and pull simulations](../how-to/push-pull.md#configure-partitions)). + +| Option | Description | +| --- | --- | +| `NAME` | Directory that partition `NAME` is mounted at, for example `data = /home/user/my_simdb_data`. | + ### `[development]` | Option | Description | diff --git a/docs/reference/server-configuration.md b/docs/reference/server-configuration.md index c656d1f2..3f0ae5f4 100644 --- a/docs/reference/server-configuration.md +++ b/docs/reference/server-configuration.md @@ -43,6 +43,7 @@ See [Set up PostgreSQL](../how-to/operate-server/set-up-postgresql.md). | `copy_files` | No | `True`/`False`: copy uploaded data files into the server's storage. Defaults to `True`. | | `copy_ids` | No | `True`/`False`: copy uploaded IMAS IDS data into the server's storage. Defaults to `True`. | | `user_upload_folder` | No | Optional staging directory clients upload into before ingest (returned by the `staging_dir` endpoint). Falls back to `upload_folder` if unset. | +| `max_append_size` | No | Maximum size in bytes of a single resumable-upload chunk, advertised to clients via the `Upload-Limit` header. Defaults to `8388608` (8 MiB). Lower it to fit a reverse proxy's request body limit. | ## `[flask]` @@ -174,6 +175,7 @@ Used by the optional | Option | Required | Description | | --- | --- | --- | | `data` | No | Directory used for partitioned data, for example `/data/simdb/partition`. | +| `http` | For `push` | Directory where resumable HTTP uploads are staged before ingestion, for example `/var/lib/simdb/http-staging`. Required for `simdb simulation push` against a v1.3 server. | ## `[role "NAME"]` diff --git a/docs/reference/uri-schemes.md b/docs/reference/uri-schemes.md index 3d7fed3a..11841b3a 100644 --- a/docs/reference/uri-schemes.md +++ b/docs/reference/uri-schemes.md @@ -1,8 +1,10 @@ # URI schemes The `inputs` and `outputs` in a [manifest](manifest-format.md) reference data -through URIs. SimDB understands two schemes: `file` for ordinary files and -`imas` for IMAS data entries. +through URIs. SimDB understands two schemes you can write yourself: `file` for +ordinary files and `imas` for IMAS data entries. Two further forms — a remote +`imas` URI and an `http` staging URI — are produced by SimDB itself and are +described at the end of this page. ## `file` scheme @@ -64,3 +66,18 @@ data. MDSplus data must have been written with Access Layer 5 (AL5) or later; Access Layer 4 (AL4) data must be migrated first. See [Migrate AL4 MDSplus data](../how-to/migrate-al4-mdsplus.md). ``` + +## `http` scheme (internal) + +While `simdb simulation push` uploads a simulation to a v1.3 server, each file +is referenced by a staging URI of the form + +``` +http:////file/ +``` + +The server resolves it against its `http` partition, copies the file into the +simulation's upload folder, and removes the staged copy (see +[Stage uploads on the server](../how-to/push-pull.md#stage-uploads-on-the-server)). +Like the remote `imas` form above, this scheme is generated by SimDB during a +push; you cannot use it in a manifest. diff --git a/pyproject.toml b/pyproject.toml index 7b747abc..c30d64c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "distro>=1.8.0", "email-validator>=1.1", "imas-python>=2.0.1", - "netCDF4>=1.5", + "netCDF4>=1.7.2", "numpy>=1.14", "pydantic>=2.10.6", "python-dateutil>=2.6", diff --git a/scripts/test_v13_ingestion.py b/scripts/test_v13_ingestion.py index 630a1dfc..cca60518 100755 --- a/scripts/test_v13_ingestion.py +++ b/scripts/test_v13_ingestion.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 """Test script for v1.3 simulation ingestion against a running server.""" -from simdb.workers.tasks import _calculate_checksum +from simdb.checksum import calculate_checksum from pathlib import Path import base64 @@ -37,7 +37,7 @@ def generate_simulation_file(): - checksum = _calculate_checksum(Path("tmp/partition_data/subdir/test_file.txt")) + checksum = calculate_checksum(Path("tmp/partition_data/subdir/test_file.txt")) return FileData( type="FILE", uri="data:///subdir/test_file.txt", @@ -46,7 +46,7 @@ def generate_simulation_file(): ) def generate_imas_file(relative_path): - checksum = _calculate_checksum(Path(f"tmp/partition_data/subdir/{relative_path}")) + checksum = calculate_checksum(Path(f"tmp/partition_data/subdir/{relative_path}")) return FileData( type="IMAS", uri=f"data:///subdir/{relative_path}", diff --git a/src/simdb/checksum.py b/src/simdb/checksum.py index 99aef230..91867922 100644 --- a/src/simdb/checksum.py +++ b/src/simdb/checksum.py @@ -1,14 +1,42 @@ import hashlib from pathlib import Path +from typing import Callable, Optional from simdb.imas.utils import SimDBUrl +#: Algorithm used for all catalog checksums. +CHECKSUM_ALGORITHM = "sha1" +#: Buffer size for reading files while hashing. Larger reads mean far fewer +#: syscalls on big files, which noticeably speeds up checksumming. +READ_CHUNK_SIZE = 1024 * 1024 -def sha1_checksum(uri: SimDBUrl) -> str: - """Generate a SHA1 checksum from the given file. + +def hash_file( + path: Path, + algorithm: str = CHECKSUM_ALGORITHM, + progress: Optional[Callable[[int], None]] = None, +) -> str: + """Return the hex digest of ``path`` computed with ``algorithm``. + + @param progress: optional callback invoked with the number of bytes read for + each block, suitable for advancing a progress bar. + """ + digest = hashlib.new(algorithm) + with path.open("rb") as file: + for chunk in iter(lambda: file.read(READ_CHUNK_SIZE), b""): + digest.update(chunk) + if progress is not None: + progress(len(chunk)) + return digest.hexdigest() + + +def file_checksum(uri: SimDBUrl, algorithm: str = CHECKSUM_ALGORITHM) -> str: + """Generate a checksum for the file at ``uri``. + + Checksums use :data:`CHECKSUM_ALGORITHM` (SHA-1). :param uri: the URI of the file to checksum - :return: a string containing the hex representation of the computed SHA1 checksum + :return: a string containing the hex representation of the computed checksum """ if uri.scheme != "file": raise ValueError(f"invalid scheme for file checksum: {uri.scheme}") @@ -21,8 +49,4 @@ def sha1_checksum(uri: SimDBUrl) -> str: if not path.is_file(): raise ValueError("File appears to be a directory") - sha1 = hashlib.sha1() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(4096), b""): - sha1.update(chunk) - return sha1.hexdigest() + return hash_file(path, algorithm) diff --git a/src/simdb/cli/commands/simulation.py b/src/simdb/cli/commands/simulation.py index 59bd0bdd..e08a43e4 100644 --- a/src/simdb/cli/commands/simulation.py +++ b/src/simdb/cli/commands/simulation.py @@ -1,19 +1,21 @@ import contextlib import sys +import time import urllib.parse from itertools import chain from pathlib import Path -from typing import Any, List, Optional, Tuple, Type +from typing import Any, List, Optional, Tuple import appdirs import click from rich.prompt import Confirm from simdb.cli.manifest import Manifest -from simdb.cli.remote_api import RemoteAPI, RemoteError +from simdb.cli.remote_api import APIError, RemoteAPI, RemoteError from simdb.config.config import Config from simdb.database import DatabaseError, get_local_db from simdb.database.models import Simulation +from simdb.enums import IngestionStatus from simdb.query import QueryType, parse_query_arg from simdb.validation import ValidationError, Validator @@ -200,19 +202,130 @@ def simulation_ingest(config: Config, manifest_file: str, alias: str): click.echo("ALIAS: " + simulation.alias + "\nUUID: " + str(simulation.uuid)) -def n_required_args_adaptor(n) -> Type[click.Command]: - class NRequiredArgs(click.Command): - NArgs = n +class OptionalRemoteCommand(click.Command): + """A command declared as `[REMOTE] ARG...` whose REMOTE may be left out.""" - def parse_args(self, ctx, args): - if len(args) == self.NArgs: - args.insert(0, "") - super().parse_args(ctx, args) + def parse_args(self, ctx, args): + arguments = [p for p in self.get_params(ctx) if isinstance(p, click.Argument)] + if self._count_values_given(ctx, args, arguments) < len(arguments): + args = ["", *args] + super().parse_args(ctx, args) - return NRequiredArgs + def _count_values_given(self, ctx, args, arguments) -> int: + """Count how many of the ARGUMENTS the command line provides a value for.""" + values = self.make_parser(ctx).parse_args(list(args))[0] + return sum(1 for argument in arguments if values.get(argument.name) is not None) -@simulation.command("push", cls=n_required_args_adaptor(1)) +def _prepare_simulation( + config: Config, api: RemoteAPI, sim_id: str, replaces: Optional[str] +) -> Simulation: + """Look up the local simulation SIM_ID and validate it against the remote + schemas.""" + db = get_local_db(config) + + simulation = db.get_simulation(sim_id) + if simulation is None: + raise click.ClickException(f"Failed to find simulation: {sim_id}") + + if replaces: + simulation.set_meta("replaces", replaces) + + schemas = api.get_validation_schemas() + try: + for schema in schemas: + Validator(schema).validate(simulation) + except ValidationError as err: + raise click.ClickException(f"Simulation does not validate: {err}") from err + + return simulation + + +def _wait_for_ingestion(api: RemoteAPI, sim_id: str, timeout: float) -> IngestionStatus: + """Poll the remote until the ingestion of SIM_ID reaches a terminal state. + + Reports every status change and returns the terminal status. + + :raise click.ClickException: if the remote cannot be reached, reports an + unknown status, or TIMEOUT seconds pass before + the ingestion finishes. + """ + max_consecutive_failures = 5 + poll_interval = 1.0 + + click.echo("Waiting for ingestion to complete...", nl=False) + last_status = None + consecutive_failures = 0 + deadline = time.monotonic() + timeout + while True: + try: + raw_status = api.get_ingestion_status(sim_id) + except RemoteError as err: + # The remote rejected the request, so retrying will not help. + click.echo() + raise click.ClickException( + f"Failed to check ingestion status: {err}" + ) from err + except ValueError as err: + # The remote reported a status this client does not know about, so + # waiting for it to change will not help. + click.echo() + raise click.ClickException( + f"Remote reported an unknown ingestion status: {err}" + ) from err + except Exception as err: + # Tolerate transient errors: the ingestion continues server-side + consecutive_failures += 1 + if consecutive_failures >= max_consecutive_failures: + click.echo() + raise click.ClickException( + f"Failed to check ingestion status " + f"{consecutive_failures} times in a row: {err}" + ) from err + if time.monotonic() >= deadline: + click.echo() + raise click.ClickException( + f"Timed out after {timeout:g}s waiting for ingestion to " + f"complete (last status: {last_status})" + ) from err + time.sleep(poll_interval) + continue + consecutive_failures = 0 + + try: + status = IngestionStatus(raw_status) + except ValueError as err: + click.echo() + raise click.ClickException( + f"Remote reported an unknown ingestion status: {raw_status}" + ) from err + + if status != last_status: + if last_status is not None: + click.echo(f" -> {status.value}", nl=False) + else: + click.echo(f" {status.value}", nl=False) + last_status = status + + if status.is_terminal(): + click.echo() + return status + + if time.monotonic() >= deadline: + click.echo() + raise click.ClickException( + f"Timed out after {timeout:g}s waiting for ingestion to complete " + f"(last status: {status.value})" + ) + + time.sleep(poll_interval) + + +@simulation.command( + "push_local", + cls=OptionalRemoteCommand, + short_help="Register a simulation whose files the REMOTE copies itself.", +) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") @@ -224,7 +337,14 @@ def parse_args(self, ctx, args): is_flag=True, help="Add the current user as a watcher of the simulation.", ) -def simulation_push( +@click.option( + "--timeout", + type=float, + default=600.0, + show_default=True, + help="Maximum number of seconds to wait for ingestion to complete.", +) +def simulation_push_local( config: Config, remote: Optional[str], sim_id: str, @@ -232,32 +352,73 @@ def simulation_push( password: Optional[str], replaces: Optional[str], add_watcher: bool, + timeout: float, ): - """Push the simulation with the given SIM_ID (UUID or alias) to the REMOTE.""" + """Register the simulation with the given SIM_ID (UUID or alias) on the REMOTE. + + Only the metadata is sent: the REMOTE copies the simulation files itself from + the recorded paths, which must therefore be reachable from the remote. + Waits for the remote ingestion to reach a terminal state, or until --timeout + seconds have passed. + """ api = RemoteAPI(remote, username, password, config) - db = get_local_db(config) + simulation = _prepare_simulation(config, api, sim_id, replaces) - simulation = db.get_simulation(sim_id) - if simulation is None: - raise click.ClickException(f"Failed to find simulation: {sim_id}") + api.push_local_simulation(simulation, add_watcher=add_watcher) - if replaces: - simulation.set_meta("replaces", replaces) + status = _wait_for_ingestion(api, simulation.uuid.hex, timeout) + if status is not IngestionStatus.COMPLETED: + raise click.ClickException( + f"Simulation ingestion failed with status: {status.value}" + ) - schemas = api.get_validation_schemas() - try: - for schema in schemas: - Validator(schema).validate(simulation) - except ValidationError as err: - raise click.ClickException(f"Simulation does not validate: {err}") from err + click.echo(f"Successfully pushed simulation {simulation.uuid}") + + +@simulation.command( + "push", + cls=OptionalRemoteCommand, + short_help="Upload a simulation and its files to the REMOTE.", +) +@pass_config +@click.argument("remote", required=False) +@click.argument("sim_id") +@click.option("--username", help="Username used to authenticate with the remote.") +@click.option("--password", help="Password used to authenticate with the remote.") +@click.option("--replaces", help="SIM_ID of simulation to deprecate and replace.") +@click.option( + "--add-watcher", + is_flag=True, + help="Add the current user as a watcher of the simulation.", +) +def simulation_push( + config: Config, + remote: Optional[str], + sim_id: str, + username: Optional[str], + password: Optional[str], + replaces: Optional[str], + add_watcher: bool, +): + """Push the simulation with the given SIM_ID (UUID or alias) to the REMOTE. - api.push_simulation(simulation, out_stream=sys.stdout, add_watcher=add_watcher) + Both the metadata and the simulation files are uploaded over HTTP. Use + push_local instead when the REMOTE can read the files itself. + """ + + api = RemoteAPI(remote, username, password, config) + simulation = _prepare_simulation(config, api, sim_id, replaces) + + try: + api.push_simulation(simulation, add_watcher=add_watcher) + except APIError as err: + raise click.ClickException(f"Failed to push simulation: {err}") from err click.echo(f"Successfully pushed simulation {simulation.uuid}") -@simulation.command("pull", cls=n_required_args_adaptor(2)) +@simulation.command("pull", cls=OptionalRemoteCommand) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") @@ -380,7 +541,7 @@ def simulation_query( ) -@simulation.command("data", cls=n_required_args_adaptor(2)) +@simulation.command("data", cls=OptionalRemoteCommand) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") @@ -448,7 +609,7 @@ def simulation_data( print_quantity(coord, label=f"coord {coord['name']}", show_stats=False) -@simulation.command("validate", cls=n_required_args_adaptor(1)) +@simulation.command("validate", cls=OptionalRemoteCommand) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") diff --git a/src/simdb/cli/remote_api.py b/src/simdb/cli/remote_api.py index 97348cdb..6ce53ffa 100644 --- a/src/simdb/cli/remote_api.py +++ b/src/simdb/cli/remote_api.py @@ -9,6 +9,7 @@ import pickle import shutil import sys +import time import uuid from collections import defaultdict from io import BytesIO @@ -25,19 +26,32 @@ Tuple, Union, ) -from urllib.parse import urlparse +from urllib.parse import quote, urlparse import appdirs import click import requests +from netCDF4 import Dataset from requests.auth import AuthBase +from rich.progress import ( + BarColumn, + DownloadColumn, + Progress, + TextColumn, + TimeRemainingColumn, + TransferSpeedColumn, +) from semantic_version import Version +from simdb.checksum import CHECKSUM_ALGORITHM, READ_CHUNK_SIZE, hash_file +from simdb.cli.resumable_upload import resumable_upload from simdb.config import Config from simdb.database.models import Simulation -from simdb.imas.utils import SimDBUrl, imas_files +from simdb.enums import IngestionStatus +from simdb.imas.utils import SimDBUrl, imas_backend_for_directory, imas_files from simdb.json import CustomDecoder, CustomEncoder from simdb.remote import CLIENT_API_VERSIONS, APIConstants +from simdb.remote.models import FileData, SimulationPostData from .manifest import DataType @@ -235,6 +249,207 @@ def _get_paths(file: "File") -> Iterable[Path]: return imas_files(file.uri) +def _check_file_is_imas(file: Path) -> bool: + # NetCDF is identified by the IMAS "Conventions" attribute + if file.suffix == ".nc": + try: + with Dataset(file, "r") as ds: + if getattr(ds, "Conventions", None) == "IMAS": + return True + except OSError: + # Not a readable NetCDF file; fall back to the directory heuristics + pass + + try: + imas_backend_for_directory(file.parent) + except ValueError: + return False + return True + + +def _partition_roots(config: Config) -> dict[str, str]: + section = config.get_section("partition", default={}) + return {k: str(v) for k, v in section.items()} + + +def _find_partition_for_file( + file: Path, partitions: Dict[str, str] +) -> Tuple[str, Path]: + # Match the partition with the longest root so that a catch-all mapping + # (e.g. "/") does not shadow more specific partitions. + best: Optional[Tuple[str, Path]] = None + best_depth = -1 + for partition, path in partitions.items(): + root = Path(path) + try: + relative = file.relative_to(root) + except ValueError: + continue + depth = len(root.parts) + if depth > best_depth: + best = (partition, relative) + best_depth = depth + if best is None: + raise APIError( + f"File {file} is not located under any configured partition " + f"(configured partitions: {', '.join(partitions) or 'none'})" + ) + return best + + +def _file_data_for_partition( + file: FileData, source: Path, partitions: Dict[str, str], keep_uuid: bool = True +) -> FileData: + """Copy FILE with its URI rewritten relative to the partition holding SOURCE.""" + partition, partition_path = _find_partition_for_file(source, partitions) + new_uri = SimDBUrl.build(scheme=partition, path=partition_path.as_posix()) + update: Dict[str, Any] = { + "uri": new_uri.encoded_string(), + "checksum": hash_file(source), + } + if not keep_uuid: + update["uuid"] = uuid.uuid1() + return file.model_copy(update=update) + + +def _source_files(file: FileData) -> List[Path]: + """Return the local files that FILE refers to, expanding directories.""" + file_uri = SimDBUrl(file.uri) + if file_uri.path is None: + raise APIError(f"File URI has no path: {file.uri}") + + if file_uri.scheme == "imas": + try: + sources = sorted(imas_files(file_uri)) + except ValueError as err: + raise APIError(f"Failed to list IMAS files of {file.uri}: {err}") from err + if not sources: + raise APIError(f"IMAS URI does not contain any files: {file.uri}") + return sources + + file_path = Path(file_uri.path) + if not file_path.is_dir(): + return [file_path] + + sources = [] + for sub_file in sorted(file_path.iterdir()): + if sub_file.is_dir(): + raise APIError(f"Nested directory found in {file_path}: {sub_file.name}") + sources.append(sub_file) + return sources + + +def _expand_directories( + files: Iterable[FileData], partitions: Dict[str, str] +) -> List[FileData]: + new_file_list = [] + for file in files: + sources = _source_files(file) + keep_uuid = len(sources) == 1 + for source in sources: + new_file_list.append( + _file_data_for_partition(file, source, partitions, keep_uuid=keep_uuid) + ) + return new_file_list + + +def _expand_directories_http( + files: Iterable[FileData], sim_uuid: uuid.UUID +) -> List[Tuple[FileData, Path, str]]: + """Expand directories / IMAS data into individual files for HTTP upload. + + Returns ``(file_data, local_source_path, target)`` triples. Structure + handling (IMAS directories stay grouped, standalone files stay flat) is + identical to local push, but unlike ``push_local`` the file bytes are + uploaded, so partitions play no role: each file keeps its absolute local + path, namespaced under ``/file/``, and is assigned an ``http://`` + URI so the server stages it into the ``http`` partition. The server's + existing copy step strips the common root exactly as it does for local + push. + """ + result: List[Tuple[FileData, Path, str]] = [] + for file in files: + file_uri = SimDBUrl(file.uri) + if file_uri.path is None: + raise ValueError("File has no associated path") + file_path = Path(file_uri.path) + if file_uri.scheme == "imas": + qs = dict(file_uri.query_params()) + path = qs.get("path") + if path is None: + raise ValueError("IMAS uri has not path set") + file_path = Path(path) + + if file_path.is_dir(): + for sub_file in file_path.iterdir(): + if sub_file.is_dir(): + raise ValueError("Nested directory found") + result.append(_make_http_entry(file, sub_file, sim_uuid)) + else: + result.append(_make_http_entry(file, file_path, sim_uuid)) + return result + + +def _make_http_entry( + template: FileData, + local_path: Path, + sim_uuid: uuid.UUID, +) -> Tuple[FileData, Path, str]: + """Build the HTTP upload entry for a single local file. + + HTTP uploads carry the file bytes from the local system, so partitions are + not consulted: the file's absolute path is namespaced under + ``/file/``, which keeps targets unique on the server. + + The checksum is left empty here and filled in later by + :func:`_compute_checksums`, so that hashing (a full read of every file) can be + reported with a progress bar instead of stalling silently before the upload. + """ + rel_posix = local_path.as_posix().lstrip("/") + target = f"{sim_uuid.hex}/file/{rel_posix}" + new_uri = SimDBUrl.build(scheme="http", path=target, host="") + file_type = "IMAS" if _check_file_is_imas(local_path) else template.type + return ( + FileData( + type=file_type, + uri=new_uri.encoded_string(), + checksum="", + datetime=template.datetime, + usage=template.usage, + purpose=template.purpose, + sensitivity=template.sensitivity, + access=template.access, + embargo=template.embargo, + ), + local_path, + target, + ) + + +def _compute_checksums(files: List[Tuple[FileData, Path, str]]) -> None: + """Compute and store the SHA-1 checksum of each file, reporting progress. + + Hashing reads every file in full and is the main delay before the upload + starts, so surface it with a byte-level progress bar (mirroring the upload + bars). The computed checksum is stored as the catalog checksum. + """ + total_bytes = sum(local_path.stat().st_size for _, local_path, _ in files) + with Progress( + TextColumn("[bold blue]{task.description}"), + BarColumn(), + DownloadColumn(), + TransferSpeedColumn(), + TimeRemainingColumn(), + ) as progress: + task = progress.add_task("Calculating checksums", total=total_bytes) + for file_data, local_path, _target in files: + progress.update(task, description=f"Hashing {local_path.name}") + file_data.checksum = hash_file( + local_path, progress=lambda n: progress.advance(task, n) + ) + progress.update(task, description="Calculated checksums") + + class RemoteAPI: """ Class to represent connection to remote API. @@ -790,16 +1005,15 @@ def _push_file( file_type: str, sim_data: Dict[str, Any], chunk_size: int, - out_stream: IO, type: DataType, ): msg = f"Uploading file {path} " - print(msg, file=out_stream, end="") + print(msg, end="") num_chunks = 0 for chunk_index, chunk in enumerate( _read_bytes_in_chunks(path, compressed=True, chunk_size=chunk_size) ): - print(".", file=out_stream, end="", flush=True) + print(".", end="") self._send_chunk(chunk_index, chunk, chunk_size, uuid, file_type, sim_data) num_chunks += 1 if num_chunks == 0: @@ -821,11 +1035,9 @@ def _push_file( ], }, ) - print(f"\r{msg}", file=out_stream, end="") + print(f"\r{msg}", end="") print( "Complete".rjust(shutil.get_terminal_size().columns - len(msg)), - file=out_stream, - flush=True, ) def _send_chunk( @@ -855,12 +1067,139 @@ def _send_chunk( ] self.post("files", data={}, files=files) - @versioned_method("v1.2", "v1.3") + @versioned_method("v1.3") + @try_request + def push_local_simulation(self, simulation: Simulation, add_watcher: bool = False): + sim_data = simulation.to_model(recurse=True) + + partitions = _partition_roots(self._config) + sim_data.inputs.root = _expand_directories(sim_data.inputs.root, partitions) + sim_data.outputs.root = _expand_directories(sim_data.outputs.root, partitions) + + uploaded_by = simulation.meta_dict().get("uploaded_by") + + post_data = SimulationPostData( + simulation=sim_data, + add_watcher=add_watcher, + uploaded_by=str(uploaded_by) if uploaded_by is not None else None, + ) + self.post("simulations", data=post_data.model_dump(mode="json")) + + def _upload_files( + self, + files: List[Tuple[FileData, Path, str]], + upload_headers: Dict[str, str], + ): + """Upload the expanded files over resumable HTTP, showing two progress + bars: an overall bar across all bytes and a sub-bar for the current file. + """ + total_bytes = sum(local_path.stat().st_size for _, local_path, _ in files) + + with Progress( + TextColumn("[bold blue]{task.description}"), + BarColumn(), + DownloadColumn(), + TransferSpeedColumn(), + TimeRemainingColumn(), + ) as progress: + overall = progress.add_task("Overall", total=total_bytes) + file_task = progress.add_task("", total=0) + uploaded = 0 + for _file_data, local_path, target in files: + size = local_path.stat().st_size + progress.reset( + file_task, total=size, description=f" {local_path.name}" + ) + url = f"{self._api_url}upload/{quote(target)}" + + def _on_progress(completed: int, _base: int = uploaded) -> None: + progress.update(file_task, completed=completed) + progress.update(overall, completed=_base + completed) + + resumable_upload( + url, + local_path, + auth=self._get_auth() if self._server_auth != "None" else None, + cookies=self._cookies, + headers=upload_headers, + progress=_on_progress, + ) + uploaded += size + progress.update(file_task, completed=size) + progress.update(overall, completed=uploaded) + + @versioned_method("v1.3") @try_request def push_simulation( + self, + simulation: Simulation, + add_watcher: bool = False, + ): + """Push a simulation by uploading its files over resumable HTTP. + + Unlike :meth:`push_local_simulation` (which requires a filesystem shared + with the server), this uploads the file bytes to the server's ``http`` + partition using a resumable protocol, then pushes the metadata. + """ + sim_data = simulation.to_model(recurse=True) + + inputs = _expand_directories_http(sim_data.inputs.root, simulation.uuid) + outputs = _expand_directories_http(sim_data.outputs.root, simulation.uuid) + + files = list(itertools.chain(inputs, outputs)) + upload_headers = {"User-Agent": "it_script_basic"} + if files: + _compute_checksums(files) + self._upload_files(files, upload_headers) + + sim_data.inputs.root = [file_data for file_data, _, _ in inputs] + sim_data.outputs.root = [file_data for file_data, _, _ in outputs] + + uploaded_by = simulation.meta_dict().get("uploaded_by") + + post_data = SimulationPostData( + simulation=sim_data, + add_watcher=add_watcher, + uploaded_by=str(uploaded_by) if uploaded_by is not None else None, + ) + self.post("simulations", data=post_data.model_dump(mode="json")) + + print("Waiting for ingestion to complete...", end="") + last_status = None + while True: + try: + status = self.get_ingestion_status(simulation.uuid.hex) + except Exception as err: + raise APIError(f"Failed to check ingestion status: {err}") from err + + if status != last_status: + if last_status is not None: + print(f" -> {status.value}", end="") + else: + print(f" {status.value}", end="") + last_status = status + + if status.is_terminal(): + break + + time.sleep(1) + + if status == IngestionStatus.COMPLETED: + return + else: + raise APIError(f"Simulation ingestion failed with status: {status.value}") + + @versioned_method("v1.3") + @try_request + def get_ingestion_status(self, sim_id: str) -> IngestionStatus: + res = self.get(f"simulation/status/{sim_id}") + return IngestionStatus(res.json()["status"]) + + @push_simulation.register("v1.2") + @try_request + def _push_simulation_v12( self, simulation: "Simulation", - out_stream: IO[str] = sys.stdout, add_watcher: bool = True, ) -> None: """ @@ -870,7 +1209,6 @@ def push_simulation( simulation metadata. :param simulation: The Simulation to push to remote server - :param out_stream: The IO stream to write messages to the user (default: stdout) :param add_watcher: Add the current user as a watcher of the simulation on the remote server """ @@ -905,7 +1243,7 @@ def push_simulation( for file in simulation.inputs: if file.type == DataType.IMAS: if not copy_ids: - print(f"Skipping IDS data {file}", file=out_stream, flush=True) + print(f"Skipping IDS data {file}") continue ids_list = simulation.meta_dict().get("input_ids", []) for path in imas_files(file.uri): @@ -927,7 +1265,6 @@ def push_simulation( "input", sim_data, chunk_size, - out_stream, file.type, ) @@ -954,14 +1291,13 @@ def push_simulation( "input", sim_data, chunk_size, - out_stream, file.type, ) for file in simulation.outputs: if file.type == DataType.IMAS: if not copy_ids: - print(f"Skipping IDS data {file}", file=out_stream, flush=True) + print(f"Skipping IDS data {file}") continue ids_list = simulation.meta_dict().get("ids", []) @@ -990,7 +1326,6 @@ def push_simulation( "output", sim_data, chunk_size, - out_stream, file.type, ) @@ -1016,13 +1351,12 @@ def push_simulation( "output", sim_data, chunk_size, - out_stream, file.type, ) sim_data = simulation.data(recurse=True) uploaded_by = simulation.meta_dict().get("uploaded_by", None) - print("Uploading simulation data ... ", file=out_stream, end="", flush=True) + print("Uploading simulation data ... ", end="") self.post( "simulations", data={ @@ -1031,7 +1365,7 @@ def push_simulation( "uploaded_by": uploaded_by, }, ) - print("Success", file=out_stream, flush=True) + print("Success") def _get_file_info(self, uuid: uuid.UUID) -> List[Tuple[Path, str]]: r = self.get(f"file/{uuid.hex}") @@ -1057,7 +1391,7 @@ def _pull_file( response = self.get(f"file/download/{uuid.hex}/{index}", stream=True) to_path.parent.mkdir(parents=True, exist_ok=True) - sha1 = hashlib.sha1() + digest = hashlib.new(CHECKSUM_ALGORITHM) with to_path.open("wb") as f: total_length = response.headers.get("content-length") @@ -1066,8 +1400,8 @@ def _pull_file( else: downloaded = 0 total_length = int(total_length) - for data in response.iter_content(chunk_size=4096): - sha1.update(data) + for data in response.iter_content(chunk_size=READ_CHUNK_SIZE): + digest.update(data) downloaded += len(data) f.write(data) done = int(50 * downloaded / total_length) @@ -1083,7 +1417,7 @@ def _pull_file( ) print("\r", file=out_stream, end="", flush=True) - if sha1.hexdigest() != checksum: + if digest.hexdigest() != checksum: raise APIError(f"Checksum failed for file {from_path}") @versioned_method("v1.2", "v1.3") diff --git a/src/simdb/cli/resumable_upload.py b/src/simdb/cli/resumable_upload.py new file mode 100644 index 00000000..d72f9c50 --- /dev/null +++ b/src/simdb/cli/resumable_upload.py @@ -0,0 +1,296 @@ +"""Client for the IETF "Resumable Uploads for HTTP" protocol. + +This is a small, dependency-free (uses ``requests``, already a dependency) +implementation of draft-ietf-httpbis-resumable-upload-11 (interop version 8) + +The single public entry point :func:`resumable_upload` uploads a local file to a +server endpoint that speaks the same protocol. The upload resource is identified +by the request URL itself: an interrupted upload can be resumed by simply +re-invoking :func:`resumable_upload` with the same arguments - the client asks +the server (via ``HEAD``) how many bytes it already has and continues from there. +""" + +import base64 +import hashlib +import logging +from pathlib import Path +from typing import Callable, Mapping, Optional, Tuple, Union + +import requests +from requests.auth import AuthBase +from requests.cookies import RequestsCookieJar + +logger = logging.getLogger(__name__) + +#: The draft interop version this client implements. +INTEROP_VERSION = "8" +INTEROP_HEADER = "Upload-Draft-Interop-Version" +#: Content type used for the body of append (``PATCH``) requests. +PARTIAL_UPLOAD_CONTENT_TYPE = "application/partial-upload" +DIGEST_ALGORITHM = "sha-256" +_HASHLIB_NAME = "sha256" +#: Default size of a single ``PATCH`` chunk (kept below the 10 MB request cap +#: enforced on the ITER network, see ``RemoteAPI.push_simulation``). +DEFAULT_CHUNK_SIZE = 8 * 1024 * 1024 + +#: Maximum number of consecutive failures (connection errors / offset +#: conflicts) tolerated before giving up. +_MAX_RETRIES = 5 + + +class ResumableUploadError(RuntimeError): + """Raised when a resumable upload cannot be completed.""" + + +def _bool_field(value: bool) -> str: + """Render a boolean as an HTTP structured-field item (``?1``/``?0``).""" + return "?1" if value else "?0" + + +def _parse_bool_field(value: Optional[str]) -> Optional[bool]: + if value is None: + return None + value = value.strip() + if value == "?1": + return True + if value == "?0": + return False + return None + + +def _format_digest(digest: bytes) -> str: + """Render a raw digest as an RFC 9530 structured-field dictionary value. + + The single member uses :data:`DIGEST_ALGORITHM` as its key and the digest as + a base64-encoded byte sequence, e.g. ``sha-256=:47DEQpj8HBSa...:``. + """ + encoded = base64.b64encode(digest).decode("ascii") + return f"{DIGEST_ALGORITHM}=:{encoded}:" + + +def _content_digest(data: bytes) -> str: + """``Content-Digest`` value for the bytes of a single request body.""" + return _format_digest(hashlib.new(_HASHLIB_NAME, data).digest()) + + +def _header_int(resp: "requests.Response", name: str) -> Optional[int]: + raw = resp.headers.get(name) + if raw is None: + return None + try: + return int(raw.strip()) + except ValueError: + return None + + +def _base_headers(extra: Optional[Mapping[str, str]]) -> dict: + headers = {INTEROP_HEADER: INTEROP_VERSION} + if extra: + headers.update(extra) + return headers + + +def _parse_upload_limit(value: Optional[str]) -> dict: + """Parse an ``Upload-Limit`` structured-field dictionary into a dict. + + Only the integer-valued members this client cares about are kept (e.g. + ``max-append-size``). Unparseable members are ignored. + """ + limits: dict = {} + if not value: + return limits + for member in value.split(","): + member = member.strip() + if "=" not in member: + continue + key, _, raw = member.partition("=") + try: + limits[key.strip()] = int(raw.strip()) + except ValueError: + continue + return limits + + +def _clamp_chunk_size(chunk_size: int, limits: dict) -> int: + """Reduce ``chunk_size`` to the server-advertised ``max-append-size``.""" + max_append = limits.get("max-append-size") + if max_append and max_append > 0: + return min(chunk_size, max_append) + return chunk_size + + +def resumable_upload( + url: str, + path: Union[str, Path], + *, + auth: Optional[Union[AuthBase, Tuple[str, str]]] = None, + cookies: Optional[Union[Mapping[str, str], RequestsCookieJar]] = None, + headers: Optional[Mapping[str, str]] = None, + chunk_size: int = DEFAULT_CHUNK_SIZE, + progress: Optional[Callable[[int], None]] = None, +) -> None: + """Upload ``path`` to ``url`` using the resumable upload protocol. + + :param url: the upload resource URL. The server is expected to treat this + URL itself as the upload resource (it is both the creation target and + the resource that is appended to / queried). + :param path: the local file to upload. + :param auth: authentication passed through to ``requests``. + :param cookies: cookies passed through to ``requests`` (e.g. firewall). + :param headers: extra headers to send with every request. + :param chunk_size: number of bytes sent per ``PATCH`` request. + :param progress: optional callback invoked with the absolute number of bytes + confirmed by the server, after resuming and after each chunk. Useful for + driving a progress bar. + """ + path = Path(path) + total = path.stat().st_size + + offset, complete, limits = _resume_or_create(url, total, auth, cookies, headers) + if complete: + if progress: + progress(total) + return + + # Reflect any bytes the server already holds (resumed upload). + if progress: + progress(offset) + + # The server advertises its append-size limit via Upload-Limit; never send a + # chunk larger than it will accept. + chunk_size = _clamp_chunk_size(chunk_size, limits) + + with path.open("rb") as f: + _send_chunks( + url, + f, + offset, + total, + chunk_size, + auth, + cookies, + headers, + progress, + ) + + +def _resume_or_create( + url: str, + total: int, + auth, + cookies, + headers, +) -> Tuple[int, bool, dict]: + """Return ``(offset, complete, limits)`` for the upload resource at ``url``. + + Probes the resource with ``HEAD``; if it does not yet exist the resource is + created with an empty body (``Upload-Complete: ?0``). ``limits`` is the + parsed ``Upload-Limit`` dictionary advertised by the server. + """ + resp = requests.head( + url, headers=_base_headers(headers), auth=auth, cookies=cookies + ) + if resp.status_code in (200, 204): + limits = _parse_upload_limit(resp.headers.get("Upload-Limit")) + if _parse_bool_field(resp.headers.get("Upload-Complete")): + return total, True, limits + return _header_int(resp, "Upload-Offset") or 0, False, limits + + create_headers = _base_headers(headers) + create_headers["Upload-Complete"] = _bool_field(False) + create_headers["Upload-Length"] = str(total) + resp = requests.post( + url, data=b"", headers=create_headers, auth=auth, cookies=cookies + ) + if resp.status_code not in (200, 201, 204): + raise ResumableUploadError( + f"Failed to create upload resource ({resp.status_code}): {resp.text}" + ) + limits = _parse_upload_limit(resp.headers.get("Upload-Limit")) + return _header_int(resp, "Upload-Offset") or 0, False, limits + + +def _send_chunks( + url: str, + f, + offset: int, + total: int, + chunk_size: int, + auth, + cookies, + headers, + progress: Optional[Callable[[int], None]] = None, +) -> None: + attempts = 0 + while True: + f.seek(offset) + chunk = f.read(chunk_size) + complete = (offset + len(chunk)) >= total + + patch_headers = _base_headers(headers) + patch_headers["Content-Type"] = PARTIAL_UPLOAD_CONTENT_TYPE + patch_headers["Upload-Offset"] = str(offset) + patch_headers["Upload-Complete"] = _bool_field(complete) + patch_headers["Content-Digest"] = _content_digest(chunk) + + try: + resp = requests.patch( + url, data=chunk, headers=patch_headers, auth=auth, cookies=cookies + ) + except (requests.ConnectionError, requests.Timeout) as err: + attempts += 1 + if attempts > _MAX_RETRIES: + raise ResumableUploadError( + f"Upload failed after {_MAX_RETRIES} retries: {err}" + ) from err + logger.warning("Upload chunk failed (%s), resuming from server offset", err) + server_offset = _query_offset(url, auth, cookies, headers) + if server_offset is not None: + offset = server_offset + continue + + if resp.status_code == 409: + # Offset mismatch: resynchronise to the server-reported offset. + server_offset = _header_int(resp, "Upload-Offset") + if server_offset is None: + raise ResumableUploadError( + "Server reported an offset conflict without an Upload-Offset header" + ) + attempts += 1 + if attempts > _MAX_RETRIES: + raise ResumableUploadError("Too many offset conflicts during upload") + offset = server_offset + continue + + if resp.status_code not in (200, 201, 204): + raise ResumableUploadError( + f"Unexpected status {resp.status_code} while appending: {resp.text}" + ) + + attempts = 0 + server_offset = _header_int(resp, "Upload-Offset") + offset = server_offset if server_offset is not None else offset + len(chunk) + + if progress: + progress(offset) + + if complete: + return + + +def _query_offset(url: str, auth, cookies, headers) -> Optional[int]: + """Return the server's current offset, or ``None`` if it can't be determined. + + A ``None`` result (failed/ambiguous HEAD, or a missing ``Upload-Offset`` + header) means "unknown" - the caller must keep its current offset rather than + restart the upload. + """ + try: + resp = requests.head( + url, headers=_base_headers(headers), auth=auth, cookies=cookies + ) + except (requests.ConnectionError, requests.Timeout): + return None + if resp.status_code in (200, 204): + return _header_int(resp, "Upload-Offset") + return None diff --git a/src/simdb/database/models/file.py b/src/simdb/database/models/file.py index 202b94dc..2bf22948 100644 --- a/src/simdb/database/models/file.py +++ b/src/simdb/database/models/file.py @@ -7,7 +7,7 @@ from sqlalchemy import Column from sqlalchemy import types as sql_types -from simdb.checksum import sha1_checksum +from simdb.checksum import file_checksum from simdb.cli.manifest import DataType from simdb.config.config import Config from simdb.docstrings import inherit_docstrings @@ -78,7 +78,7 @@ def generate_checksum(self, config, ids_list: list): elif self.type == DataType.IMAS: checksum = imas_checksum(self.uri, ids_list) elif self.type == DataType.FILE: - checksum = sha1_checksum(self.uri) + checksum = file_checksum(self.uri) else: raise NotImplementedError(f"Cannot generate checksum for type {self.type}.") return checksum @@ -139,7 +139,12 @@ def to_model_with_path(self) -> FileGetDataResponse: files = [FileInfo(path=Path(self.uri.path), checksum=self.checksum)] else: files = [ - FileInfo(path=path, checksum=sha1_checksum(SimDBUrl(f"file:{path}"))) + FileInfo( + path=path, + checksum=file_checksum( + SimDBUrl.build(scheme="file", path=path.as_posix()) + ), + ) for path in imas_files(self.uri) ] return FileGetDataResponse( diff --git a/src/simdb/imas/checksum.py b/src/simdb/imas/checksum.py index d9d403ef..52717a5c 100644 --- a/src/simdb/imas/checksum.py +++ b/src/simdb/imas/checksum.py @@ -1,6 +1,7 @@ import hashlib from pathlib import Path +from simdb.checksum import CHECKSUM_ALGORITHM, READ_CHUNK_SIZE from simdb.imas.utils import SimDBUrl from .utils import imas_files, list_idss, open_imas @@ -8,8 +9,8 @@ IGNORED_FIELDS = ("data_dictionary", "access_layer", "access_layer_language") -def checksum(uri: SimDBUrl, ids_list: list) -> str: - sha1 = hashlib.sha1() +def checksum(uri: SimDBUrl, ids_list: list, algorithm: str = CHECKSUM_ALGORITHM) -> str: + digest = hashlib.new(algorithm) if not ids_list: entry = open_imas(uri) @@ -25,6 +26,6 @@ def checksum(uri: SimDBUrl, ids_list: list) -> str: and ids_name[0] not in ids_list ): continue - for chunk in iter(lambda: file.read(4096), b""): - sha1.update(chunk) - return sha1.hexdigest() + for chunk in iter(lambda: file.read(READ_CHUNK_SIZE), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/src/simdb/imas/utils.py b/src/simdb/imas/utils.py index 73a4c73e..c54aee2d 100644 --- a/src/simdb/imas/utils.py +++ b/src/simdb/imas/utils.py @@ -35,15 +35,16 @@ def build( fragment: Optional[str] = None, **kwargs, ) -> "SimDBUrl": - url_str = f"{scheme}:" - + authority = "" if host: - url_str += f"//{host}" - if port: - url_str += f":{port}" - url_str += "/" - - url_str += path or "" + authority = f"//{host}" + if port is not None: + authority += f":{port}" + path_part = path or "" + if authority and path_part and not path_part.startswith("/"): + path_part = f"/{path_part}" + + url_str = f"{scheme}:{authority}{path_part}" if query: url_str += f"?{query}" if fragment: @@ -103,10 +104,11 @@ def list_idss(entry: DBEntry) -> List[str]: for ids_name in entry.factory.ids_names(): occurrences = entry.list_all_occurrences(ids_name) if occurrences and len(occurrences) > 0: - for occurrence in range(len(occurrences)): + for occurrence in occurrences: if occurrence > 0: idss.append(ids_name + "_" + str(occurrence)) - idss.append(ids_name) + else: + idss.append(ids_name) return idss @@ -222,10 +224,19 @@ def open_imas(uri: SimDBUrl) -> DBEntry: if uri.scheme == "file": imas_uri = uri.path elif uri.scheme == "imas": + # Access Layer 4 / legacy entries are opened through the dedicated + # DBEntry constructor rather than a URI string. + if not _is_al5(): + return _open_legacy(uri) + qs = dict(uri.query_params()) path = qs.get("path") if path is None: - raise ValueError(f"invalid imas URI: {uri} - no path found") + # A legacy-style URI (no explicit path query): resolve the on-disk + # path and rebuild an AL5 URI before opening. + path = get_path_for_legacy_uri(uri) + backend = qs.get("backend", "mdsplus") + uri = SimDBUrl.build(scheme="imas", path=backend, query=f"path={path}") imas_uri = str(uri) else: raise ValueError(f"invalid imas URI: {uri} - invalid scheme") @@ -311,6 +322,37 @@ def _get_path(uri: SimDBUrl) -> Path: return path +def imas_backend_for_directory(directory: Path) -> str: + """ + Identify the IMAS backend of a directory by inspecting its contents. + + @param directory: a directory that may contain an IMAS dataset + @return: the backend name ("ascii", "hdf5" or "mdsplus") + @raise ValueError: if no IMAS dataset is detected + """ + children = list(directory.iterdir()) + + # ASCII heuristic + if any(child.suffix == ".ids" for child in children): + return "ascii" + + # HDF5 heuristic + if any(child.suffix == ".h5" for child in children) and any( + child.name == "master.h5" for child in children + ): + return "hdf5" + + # MDSplus heuristic + if {p.name for p in children} >= { + "ids_001.tree", + "ids_001.characteristics", + "ids_001.datafile", + }: + return "mdsplus" + + raise ValueError("IMAS backend could not be identified.") + + def imas_files(uri: SimDBUrl) -> List[Path]: """ Return all the files associated with the given IMAS URI. diff --git a/src/simdb/remote/apis/files.py b/src/simdb/remote/apis/files.py index cd7cad5e..61286623 100644 --- a/src/simdb/remote/apis/files.py +++ b/src/simdb/remote/apis/files.py @@ -8,7 +8,7 @@ from flask_restx import Namespace, Resource from werkzeug.datastructures import FileStorage -from simdb.checksum import sha1_checksum +from simdb.checksum import file_checksum from simdb.cli.manifest import DataType from simdb.database import DatabaseError, models from simdb.imas.checksum import checksum as imas_checksum @@ -49,7 +49,9 @@ def _verify_file( path = secure_path(Path(sim_file.uri.path), common_root, staging_dir) if not path.exists(): raise ValueError(f"file {path} does not exist") - checksum = sha1_checksum(SimDBUrl.build(scheme="file", path=path.as_posix())) + checksum = file_checksum( + SimDBUrl.build(scheme="file", host="", path=path.as_posix()), + ) if sim_file.checksum != checksum: raise ValueError(f"checksum failed for file {sim_file!r}") elif sim_file.type == DataType.IMAS: @@ -66,7 +68,7 @@ def _verify_file( else: path_value = str(staging_dir) new_uri = uri.build( - scheme=uri.scheme, path=uri.path, query=f"path={path_value}" + scheme=uri.scheme, host="", path=uri.path, query=f"path={path_value}" ) checksum = imas_checksum(new_uri, ids_list or []) if sim_file.checksum != checksum: diff --git a/src/simdb/remote/apis/v1_3/__init__.py b/src/simdb/remote/apis/v1_3/__init__.py index a0bcf15d..67a54293 100644 --- a/src/simdb/remote/apis/v1_3/__init__.py +++ b/src/simdb/remote/apis/v1_3/__init__.py @@ -9,6 +9,7 @@ from simdb.remote.core.auth import TokenAuthenticator from .simulation_data import api as data_ns +from .upload import api as upload_ns api = Api( title="SimDB REST API", @@ -28,7 +29,7 @@ doc="/docs", ) -namespaces = [metadata_ns, watcher_ns, file_ns, sim_ns, data_ns] +namespaces = [metadata_ns, watcher_ns, file_ns, sim_ns, data_ns, upload_ns] api.route("/staging_dir", defaults={"sim_hex": None})(StagingDirectory) api.route("/staging_dir/")(StagingDirectory) diff --git a/src/simdb/remote/apis/v1_3/simulations.py b/src/simdb/remote/apis/v1_3/simulations.py index 7bb0fcb6..7611f655 100644 --- a/src/simdb/remote/apis/v1_3/simulations.py +++ b/src/simdb/remote/apis/v1_3/simulations.py @@ -7,6 +7,7 @@ from simdb.database.models import simulation as models_sim from simdb.database.models import watcher as models_watcher from simdb.enums import IngestionStatus +from simdb.imas.utils import SimDBUrl from simdb.remote.apis.v1_2.simulations import ( Simulation, SimulationMeta, @@ -30,6 +31,7 @@ SimulationStatusResponse, ) from simdb.workers.tasks import ( + cleanup_http_staging_task, complete_ingestion_task, copy_files_task, ) @@ -116,8 +118,19 @@ def post( # The complete job will set simulation.ingestion_status = Completed complete = complete_ingestion_task.si(simulation.uuid) + # Files uploaded over HTTP are staged in the ``http`` partition; once + # copied into the upload folder, remove those staged duplicates. + all_files = [*body.simulation.inputs.root, *body.simulation.outputs.root] + if all(SimDBUrl(f.uri).scheme == "http" for f in all_files): + cleanup = cleanup_http_staging_task.si(simulation.uuid) + copy_files.link_error(cleanup) + complete.link_error(cleanup) + chain = copy_files | complete | cleanup + else: + chain = copy_files | complete + try: - _ = (copy_files | complete).apply_async() + _ = chain.apply_async() except Exception as err: simulation.ingestion_status = IngestionStatus.COPY_FAILED current_app.db.session.commit() diff --git a/src/simdb/remote/apis/v1_3/upload.py b/src/simdb/remote/apis/v1_3/upload.py new file mode 100644 index 00000000..640b4987 --- /dev/null +++ b/src/simdb/remote/apis/v1_3/upload.py @@ -0,0 +1,247 @@ +"""Server side of the IETF "Resumable Uploads for HTTP" protocol. + +Implements draft-ietf-httpbis-resumable-upload-11 (interop version 8), the +counterpart to :mod:`simdb.cli.resumable_upload`. Uploaded bytes are staged into +the ``http`` partition (config ``partition.http``): a client uploading to +``/v1.3/upload//`` results in the file being written to +``//``. The simulation is then pushed +(metadata only) referencing those files with ``http:////`` +URIs, which the existing ingestion pipeline resolves via the ``http`` partition. + +Upload state lives on disk so it survives across worker processes: in-progress +bytes are written to ``.partial`` and atomically renamed to ```` +when the upload completes. The current offset is simply the size of that file. +""" + +import base64 +import binascii +import contextlib +import hashlib +from pathlib import Path +from typing import Optional, Tuple + +from flask import Response, request +from flask_restx import Namespace, Resource +from werkzeug.exceptions import Forbidden + +from simdb.remote.core.auth import User, requires_auth +from simdb.remote.core.typing import current_app + +api = Namespace("upload", path="/") + +#: The draft interop version this server implements. +INTEROP_VERSION = "8" +INTEROP_HEADER = "Upload-Draft-Interop-Version" +PARTIAL_SUFFIX = ".partial" +#: Default maximum size of a single append (``PATCH``) body, advertised to +#: clients via the ``Upload-Limit`` header. Overridable with the +#: ``server.max_append_size`` config option. +DEFAULT_MAX_APPEND_SIZE = 8 * 1024 * 1024 +_DIGEST_ALGORITHMS = {"sha-256": "sha256", "sha-512": "sha512"} + + +def _max_append_size() -> int: + value = current_app.simdb_config.get_option( + "server.max_append_size", default=DEFAULT_MAX_APPEND_SIZE + ) + try: + return int(value) + except (TypeError, ValueError): + return DEFAULT_MAX_APPEND_SIZE + + +def _read_capped_body(limit: int) -> Optional[bytes]: + """Read the request body, capped at ``limit`` bytes.""" + data = request.stream.read(limit + 1) + if len(data) > limit: + return None + return data + + +def _bool_field(value: bool) -> str: + return "?1" if value else "?0" + + +def _parse_bool_field(value: Optional[str]) -> Optional[bool]: + if value is None: + return None + value = value.strip() + if value == "?1": + return True + if value == "?0": + return False + return None + + +def _parse_digest_header(value: Optional[str]) -> dict: + """Parse an RFC 9530 digest structured-field dictionary into ``{algo: bytes}``. + + Members are ``algo=:base64:`` items; only the algorithms in + :data:`_DIGEST_ALGORITHMS` are kept. Unparseable members are skipped. Commas + are safe separators here because base64 never contains them. + """ + digests: dict = {} + if not value: + return digests + for member in value.split(","): + member = member.strip() + if "=" not in member: + continue + key, _, raw = member.partition("=") + key = key.strip().lower() + if key not in _DIGEST_ALGORITHMS: + continue + raw = raw.strip() + if len(raw) >= 2 and raw.startswith(":") and raw.endswith(":"): + raw = raw[1:-1] + try: + digests[key] = base64.b64decode(raw, validate=True) + except (binascii.Error, ValueError): + continue + return digests + + +def _digest_matches(value: Optional[str], data: bytes) -> bool: + """Return whether the digest header ``value`` matches ``data``. + + Passes when no recognised algorithm is present (nothing to verify) and when + every recognised algorithm's digest matches; fails on any mismatch. + """ + provided = _parse_digest_header(value) + return all( + hashlib.new(_DIGEST_ALGORITHMS[algo], data).digest() == expected + for algo, expected in provided.items() + ) + + +def _partition_base() -> Path: + base = current_app.simdb_config.get_string_option("partition.http", default=None) + if not base: + raise ValueError("Partition 'http' is not configured on the server") + return Path(base).resolve() + + +def _resolve_target(target: str) -> Tuple[Path, Path]: + """Resolve ``target`` to its ``(final, partial)`` paths within ``partition.http``. + + Raises ``ValueError`` if the resolved path would escape the partition. + """ + base = _partition_base() + final = (base / target).resolve() + if not final.is_relative_to(base): + raise Forbidden("Access denied.") + partial = final.parent / (final.name + PARTIAL_SUFFIX) + return final, partial + + +def _state(final: Path, partial: Path) -> Tuple[int, bool, bool]: + """Return ``(offset, complete, exists)`` for the upload resource.""" + if partial.exists(): + return partial.stat().st_size, False, True + if final.exists(): + return final.stat().st_size, True, True + return 0, False, False + + +def _headers(offset: int, complete: bool) -> dict: + return { + INTEROP_HEADER: INTEROP_VERSION, + "Upload-Offset": str(offset), + "Upload-Complete": _bool_field(complete), + # Advertise the server's append-size limit (structured-field dictionary) + # so the client sizes its chunks accordingly. + "Upload-Limit": f"max-append-size={_max_append_size()}", + "Cache-Control": "no-store", + } + + +@api.route("/upload/") +class ResumableUpload(Resource): + """A single resumable upload resource staged into the ``http`` partition.""" + + @requires_auth() + def post(self, target: str, user: User) -> Response: + """Create (or reset) the upload resource and optionally write data.""" + final, partial = _resolve_target(target) + partial.parent.mkdir(parents=True, exist_ok=True) + + data = _read_capped_body(_max_append_size()) + if data is None: + return Response(status=413, headers=_headers(0, False)) + # Per-request integrity: reject before writing anything if the body does + # not match the client's Content-Digest. + if not _digest_matches(request.headers.get("Content-Digest"), data): + return Response(status=400, headers=_headers(0, False)) + with partial.open("wb") as f: + f.write(data) + offset = len(data) + + complete = _parse_bool_field(request.headers.get("Upload-Complete")) or False + if complete: + partial.replace(final) + + headers = _headers(offset, complete) + headers["Location"] = request.url + return Response(status=201, headers=headers) + + @requires_auth() + def head(self, target: str, user: User) -> Response: + """Report the current offset / completeness of the upload resource.""" + final, partial = _resolve_target(target) + offset, complete, exists = _state(final, partial) + if not exists: + return Response(status=404, headers={INTEROP_HEADER: INTEROP_VERSION}) + return Response(status=204, headers=_headers(offset, complete)) + + @requires_auth() + def patch(self, target: str, user: User) -> Response: + """Append data to the upload resource at the given ``Upload-Offset``.""" + final, partial = _resolve_target(target) + offset, complete, _exists = _state(final, partial) + + # Appending to an already-completed upload is a no-op when the client is + # simply confirming completion at the final offset. + if complete: + return Response(status=200, headers=_headers(offset, True)) + + try: + requested_offset = int(request.headers.get("Upload-Offset", "")) + except ValueError: + return Response(status=400, headers={INTEROP_HEADER: INTEROP_VERSION}) + + if requested_offset != offset: + # Offset mismatch - tell the client our current offset so it resyncs. + return Response(status=409, headers=_headers(offset, False)) + + data = _read_capped_body(_max_append_size()) + if data is None: + return Response(status=413, headers=_headers(offset, False)) + + # Per-request integrity: reject (without appending) if the body does not + # match the client's Content-Digest. Offset is left unchanged so the + # client can safely retry the same chunk. + if not _digest_matches(request.headers.get("Content-Digest"), data): + return Response(status=400, headers=_headers(offset, False)) + + partial.parent.mkdir(parents=True, exist_ok=True) + with partial.open("ab") as f: + f.write(data) + offset += len(data) + + request_complete = ( + _parse_bool_field(request.headers.get("Upload-Complete")) or False + ) + if request_complete: + partial.replace(final) + return Response(status=200, headers=_headers(offset, True)) + + return Response(status=204, headers=_headers(offset, False)) + + @requires_auth() + def delete(self, target: str, user: User) -> Response: + """Cancel the upload and remove any staged data.""" + final, partial = _resolve_target(target) + for path in (partial, final): + with contextlib.suppress(FileNotFoundError): + path.unlink() + return Response(status=204, headers={INTEROP_HEADER: INTEROP_VERSION}) diff --git a/src/simdb/workers/tasks.py b/src/simdb/workers/tasks.py index 97d9213d..9243cddb 100644 --- a/src/simdb/workers/tasks.py +++ b/src/simdb/workers/tasks.py @@ -1,4 +1,3 @@ -import hashlib import itertools import logging import os @@ -8,14 +7,13 @@ from typing import Iterable, List from uuid import UUID -from pydantic import AnyUrl - +from simdb.checksum import hash_file from simdb.config import Config from simdb.database.database import get_db from simdb.database.models import File from simdb.email.server import EmailServer from simdb.enums import IngestionStatus -from simdb.imas.utils import SimDBUrl +from simdb.imas.utils import SimDBUrl, imas_backend_for_directory from simdb.remote.models import FileData, FileDataList from simdb.workers.celery import celery_app @@ -62,36 +60,14 @@ def _imas_path_to_uri(imas_path: Path) -> SimDBUrl: if imas_path.suffix == ".nc": return SimDBUrl.build(scheme="file", path=imas_path.as_posix()) - children = set(imas_path.iterdir()) - - if any(child.suffix == ".ids" for child in children): - u = SimDBUrl.build( - scheme="imas", path="ascii", query=f"path={imas_path.as_posix()}" - ) - return u + backend = imas_backend_for_directory(imas_path) - if any(child.suffix == ".h5" for child in children) and any( - child.name == "master.h5" for child in children - ): - u = SimDBUrl.build( - scheme="imas", path="hdf5", query=f"path={imas_path.as_posix()}" - ) - return u - - if {p.name for p in children} >= { - "ids_001.tree", - "ids_001.characteristics", - "ids_001.datafile", - }: - u = SimDBUrl.build( - scheme="imas", path="mdsplus", query=f"path={imas_path.as_posix()}" - ) - return u - - raise ValueError("IMAS backend could not be identified.") + return SimDBUrl.build( + scheme="imas", path=backend, query=f"path={imas_path.as_posix()}" + ) -def _resolve_uri_to_path(uri: AnyUrl, config: Config) -> Path: +def _resolve_uri_to_path(uri: SimDBUrl, config: Config) -> Path: partition = uri.scheme if not partition: raise ValueError("Partition not given") @@ -106,6 +82,10 @@ def _resolve_uri_to_path(uri: AnyUrl, config: Config) -> Path: raise ValueError("Path not given") path = Path(path) path = path.relative_to(path.anchor) + # Standard schemes (e.g. ``http``) parse the first path segment as the URL + # authority. Fold it back in so the full relative path is reconstructed. + if uri.host: + path = Path(uri.host) / path target = (partition_path / path).resolve() if not target.is_relative_to(partition_path): raise ValueError("Access denied.") @@ -133,12 +113,9 @@ def _copy_files( shutil.copy2(source, destination) -def _calculate_checksum(path: Path) -> str: - sha1 = hashlib.sha1() - with path.open("rb") as f: - for chunk in iter(lambda: f.read(4096), b""): - sha1.update(chunk) - return sha1.hexdigest() +def _checksum_matches(path: Path, expected: str) -> bool: + """Whether ``path`` matches ``expected``.""" + return hash_file(path) == expected def _get_imas_identifier_path(path: Path) -> Path: @@ -153,8 +130,7 @@ def _create_file_from_data( uri = SimDBUrl(data.uri) path = _resolve_uri_to_path(uri, config) - checksum = _calculate_checksum(path) - if data.checksum != checksum: + if not _checksum_matches(path, data.checksum): raise ValueError("Hash of file does not match provided checksum") file = File.from_data_model(data) @@ -182,8 +158,7 @@ def _create_files_from_data_list( seen_imas_paths.add(imas_path) file = _create_file_from_data(file_data, config, imas_path) else: - checksum = _calculate_checksum(path) - if file_data.checksum != checksum: + if not _checksum_matches(path, file_data.checksum): raise ValueError("Hash of file does not match provided checksum") file = File.from_data_model(file_data) file.uri = SimDBUrl.build(scheme="file", path=path.as_posix()) @@ -314,3 +289,25 @@ def fail_stale_ingestions_task() -> dict: return {"failed": failed} finally: database.close() + + +@celery_app.task +def cleanup_http_staging_task(simulation_uuid: UUID): + """Remove a simulation's staged files from the ``http`` partition. + + HTTP-uploaded files are staged into the ``http`` partition and then copied + into the simulation's upload folder by :func:`copy_files_task`. Once copied + they are duplicates, so the staging directory is removed here. + """ + config = Config() + config.load() + + partition_path_str = config.get_string_option("partition.http", default=None) + if not partition_path_str: + return + + partition_path = Path(partition_path_str).resolve() + staging_dir = (partition_path / simulation_uuid.hex).resolve() + # Guard against escaping the partition before removing anything. + if staging_dir.is_relative_to(partition_path) and staging_dir != partition_path: + shutil.rmtree(staging_dir, ignore_errors=True) diff --git a/tests/cli/test_cli_simulation_command.py b/tests/cli/test_cli_simulation_command.py index 0120fc02..1f81def4 100644 --- a/tests/cli/test_cli_simulation_command.py +++ b/tests/cli/test_cli_simulation_command.py @@ -1,9 +1,13 @@ from unittest import mock +from uuid import uuid1 +import pytest from click.testing import CliRunner from utils import config_test_file +from simdb.cli.remote_api import FailedConnection, RemoteError from simdb.cli.simdb import cli +from simdb.enums import IngestionStatus @mock.patch("simdb.database.get_local_db") @@ -85,3 +89,199 @@ def test_simulation_validate_command(remote_api, get_local_db): runner = CliRunner() result = runner.invoke(cli, [f"--config-file={config_file}", "simulation"]) assert result.exception is None + + +@pytest.mark.parametrize( + "subcommand, trailing_args", + ( + ("push", ()), + ("push_local", ()), + ("pull", ("directory",)), + ("data", ("ids_path",)), + ("validate", ()), + ), +) +@pytest.mark.parametrize( + "options", (("--username", "bob"), ()), ids=("username", "none") +) +@pytest.mark.parametrize( + "options_first", (True, False), ids=("options-first", "options-last") +) +@pytest.mark.parametrize("remote", (("test",), ()), ids=("test", "default")) +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_optional_remote_argument( + get_local_db, remote_api, remote, options_first, options, subcommand, trailing_args +): + """REMOTE may be left out, wherever the options appear on the command line.""" + config_file = config_test_file() + # push_local polls until the ingestion reaches a terminal state. + remote_api.return_value.get_ingestion_status.return_value = ( + IngestionStatus.COMPLETED.value + ) + arguments = (*remote, "sim_id", *trailing_args) + argv = (*options, *arguments) if options_first else (*arguments, *options) + + runner = CliRunner() + result = runner.invoke( + cli, [f"--config-file={config_file}", "simulation", subcommand, *argv] + ) + + assert remote_api.called, result.output + used_remote, used_username = remote_api.call_args.args[:2] + assert used_remote == (remote[0] if remote else "") + assert used_username == ("bob" if options else None) + + +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_command(get_local_db, remote_api): + """push_local waits for the remote ingestion to complete.""" + config_file = config_test_file() + remote_api.return_value.get_ingestion_status.return_value = ( + IngestionStatus.COMPLETED.value + ) + simulation = get_local_db.return_value.get_simulation.return_value + + runner = CliRunner() + result = runner.invoke( + cli, [f"--config-file={config_file}", "simulation", "push_local", "sim_id"] + ) + + assert result.exception is None, result.output + remote_api.return_value.push_local_simulation.assert_called_once_with( + simulation, add_watcher=False + ) + assert f"Successfully pushed simulation {simulation.uuid}" in result.output + + +@pytest.mark.parametrize( + "status", (IngestionStatus.COPY_FAILED, IngestionStatus.VALIDATION_FAILED) +) +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_failed_ingestion(get_local_db, remote_api, status): + """push_local fails when the remote ingestion does.""" + config_file = config_test_file() + remote_api.return_value.get_ingestion_status.return_value = status.value + + runner = CliRunner() + result = runner.invoke( + cli, [f"--config-file={config_file}", "simulation", "push_local", "sim_id"] + ) + + assert result.exit_code != 0 + assert f"Simulation ingestion failed with status: {status.value}" in result.output + + +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_timeout(get_local_db, remote_api): + """push_local stops waiting once the --timeout has passed.""" + config_file = config_test_file() + remote_api.return_value.get_ingestion_status.return_value = ( + IngestionStatus.COPYING.value + ) + argv = ("push_local", "sim_id", "--timeout=0") + + runner = CliRunner() + result = runner.invoke(cli, [f"--config-file={config_file}", "simulation", *argv]) + + assert result.exit_code != 0 + assert "Timed out after 0s waiting for ingestion to complete" in result.output + assert "last status: COPYING" in result.output + + +def _invoke_push_local(remote_api, *extra_args): + config_file = config_test_file() + runner = CliRunner() + return runner.invoke( + cli, + [ + f"--config-file={config_file}", + "simulation", + "push_local", + "sim_id", + *extra_args, + ], + ) + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_reports_status_changes(get_local_db, remote_api, sleep): + """push_local reports every ingestion status the remote goes through.""" + sim_uuid = uuid1() + get_local_db.return_value.get_simulation.return_value.uuid = sim_uuid + remote_api.return_value.get_ingestion_status.side_effect = [ + IngestionStatus.QUEUED.value, + IngestionStatus.COPYING.value, + IngestionStatus.COPYING.value, + IngestionStatus.COMPLETED.value, + ] + + result = _invoke_push_local(remote_api) + + assert result.exit_code == 0, result.output + assert "QUEUED -> COPYING -> COMPLETED" in result.output + assert f"Successfully pushed simulation {sim_uuid}" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_unknown_status(get_local_db, remote_api, sleep): + """push_local does not wait for a status it cannot recognise.""" + remote_api.return_value.get_ingestion_status.return_value = "TELEPORTING" + + result = _invoke_push_local(remote_api) + + assert result.exit_code != 0 + assert "Remote reported an unknown ingestion status: TELEPORTING" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_transient_errors(get_local_db, remote_api, sleep): + """The ingestion continues server-side, so a failed status check is not fatal.""" + remote_api.return_value.get_ingestion_status.side_effect = [ + FailedConnection("connection reset"), + FailedConnection("connection reset"), + IngestionStatus.COMPLETED.value, + ] + + result = _invoke_push_local(remote_api) + + assert result.exit_code == 0, result.output + assert "COMPLETED" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_persistent_errors(get_local_db, remote_api, sleep): + """push_local gives up once the status check keeps failing.""" + remote_api.return_value.get_ingestion_status.side_effect = FailedConnection("down") + + result = _invoke_push_local(remote_api) + + assert result.exit_code != 0 + assert "Failed to check ingestion status 5 times in a row: down" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_rejected_status_check(get_local_db, remote_api, sleep): + """A request the remote rejects will not start working when repeated.""" + remote_api.return_value.get_ingestion_status.side_effect = RemoteError( + "Simulation not found" + ) + + result = _invoke_push_local(remote_api) + + assert result.exit_code != 0 + assert "Failed to check ingestion status: Simulation not found" in result.output + assert remote_api.return_value.get_ingestion_status.call_count == 1 diff --git a/tests/cli/test_push_http.py b/tests/cli/test_push_http.py new file mode 100644 index 00000000..83084773 --- /dev/null +++ b/tests/cli/test_push_http.py @@ -0,0 +1,93 @@ +"""Tests for the HTTP push client helpers and CLI command.""" + +import hashlib +import uuid +from datetime import datetime, timezone + +from simdb.cli.remote_api import _compute_checksums, _expand_directories_http +from simdb.imas.utils import SimDBUrl +from simdb.remote.models import FileData + + +def _file_data(path) -> FileData: + return FileData( + type="FILE", + uri=SimDBUrl.build(scheme="file", path=str(path), host="").encoded_string(), + checksum="ignored", + datetime=datetime.now(timezone.utc), + ) + + +def test_expand_directories_http_uses_absolute_paths(tmp_path): + # HTTP uploads carry the file bytes, so partitions play no role: files keep + # their absolute local path, namespaced under /file/. + (tmp_path / "subdir").mkdir() + f = tmp_path / "subdir" / "file.txt" + f.write_text("hello") + sim_uuid = uuid.uuid4() + + result = _expand_directories_http([_file_data(f)], sim_uuid) + + assert len(result) == 1 + file_data, local_path, target = result[0] + assert local_path == f + assert target == f"{sim_uuid.hex}/file/{str(f).lstrip('/')}" + parsed = SimDBUrl(file_data.uri) + assert parsed.scheme == "http" + assert parsed.host == sim_uuid.hex + assert parsed.path == f"/file/{str(f).lstrip('/')}" + assert file_data.type == "FILE" + assert file_data.checksum == "" + + +def test_compute_checksums_populates_sha1(tmp_path): + f1 = tmp_path / "a.txt" + f1.write_bytes(b"hello") + f2 = tmp_path / "b.txt" + f2.write_bytes(b"world!!") + + fd1 = _file_data(f1) + fd2 = _file_data(f2) + files = [(fd1, f1, "a"), (fd2, f2, "b")] + + _compute_checksums(files) + + assert fd1.checksum == hashlib.sha1(b"hello").hexdigest() + assert fd2.checksum == hashlib.sha1(b"world!!").hexdigest() + + +def test_expand_directories_http_keeps_imas_directory(tmp_path): + # An IMAS (hdf5) directory must stay contained in its own folder. + imas_dir = tmp_path / "run" / "myids" + imas_dir.mkdir(parents=True) + (imas_dir / "master.h5").write_text("m") + (imas_dir / "0001.h5").write_text("d") + sim_uuid = uuid.uuid4() + + imas_file = FileData( + type="IMAS", + uri=SimDBUrl.build( + scheme="imas", path="hdf5", host="", query=f"path={imas_dir}" + ).encoded_string(), + checksum="ignored", + datetime=datetime.now(timezone.utc), + ) + + result = _expand_directories_http([imas_file], sim_uuid) + + prefix = f"{sim_uuid.hex}/file/{str(imas_dir).lstrip('/')}" + targets = sorted(t for _, _, t in result) + assert targets == [f"{prefix}/0001.h5", f"{prefix}/master.h5"] + assert all(file_data.type == "IMAS" for file_data, _, _ in result) + + +def test_expand_directories_http_unpartitioned_file_uses_file_scheme(tmp_path): + # No partition configuration is needed for HTTP uploads. + f = tmp_path / "loose.txt" + f.write_text("x") + sim_uuid = uuid.uuid4() + + result = _expand_directories_http([_file_data(f)], sim_uuid) + + _, _, target = result[0] + assert target == f"{sim_uuid.hex}/file/{str(f).lstrip('/')}" diff --git a/tests/cli/test_remote_api_push_local.py b/tests/cli/test_remote_api_push_local.py new file mode 100644 index 00000000..7bb346c4 --- /dev/null +++ b/tests/cli/test_remote_api_push_local.py @@ -0,0 +1,108 @@ +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +from simdb.checksum import hash_file +from simdb.cli.remote_api import ( + APIError, + _expand_directories, + _find_partition_for_file, +) +from simdb.remote.models import FileData + + +def _file_data(uri: str, file_type: str = "FILE") -> FileData: + return FileData( + type=file_type, + uri=uri, + checksum="stale", + datetime=datetime(2026, 1, 1, tzinfo=timezone.utc), + ) + + +def test_find_partition_prefers_the_deepest_root(): + """A catch-all partition does not shadow a more specific one.""" + partitions = {"root": "/", "data": "/mnt/data"} + + assert _find_partition_for_file(Path("/mnt/data/run/x.txt"), partitions) == ( + "data", + Path("run/x.txt"), + ) + assert _find_partition_for_file(Path("/sdcc/run/x.txt"), partitions) == ( + "root", + Path("sdcc/run/x.txt"), + ) + + +def test_find_partition_without_a_match(): + with pytest.raises(APIError, match="configured partitions: data"): + _find_partition_for_file(Path("/elsewhere/x.txt"), {"data": "/mnt/data"}) + + +def test_expand_directories_rewrites_the_uri_of_a_single_file(tmp_path: Path): + source = tmp_path / "run" / "x.txt" + source.parent.mkdir() + source.write_text("contents") + file = _file_data(f"file:{source}") + + expanded = _expand_directories([file], {"data": str(tmp_path)}) + + assert len(expanded) == 1 + assert expanded[0].uri == "data:run/x.txt" + assert expanded[0].checksum == hash_file(source) + # A file that maps onto a single source keeps its identity. + assert expanded[0].uuid == file.uuid + + +def test_expand_directories_expands_a_directory(tmp_path: Path): + directory = tmp_path / "run" + directory.mkdir() + for name in ("b.txt", "a.txt"): + (directory / name).write_text(name) + file = _file_data(f"file:{directory}") + + expanded = _expand_directories([file], {"data": str(tmp_path)}) + + assert [f.uri for f in expanded] == ["data:run/a.txt", "data:run/b.txt"] + # Each file needs its own identity to be stored on the remote. + assert len({f.uuid for f in expanded}) == 2 + + +def test_expand_directories_rejects_nested_directories(tmp_path: Path): + directory = tmp_path / "run" + (directory / "nested").mkdir(parents=True) + file = _file_data(f"file:{directory}") + + with pytest.raises(APIError, match="Nested directory found"): + _expand_directories([file], {"data": str(tmp_path)}) + + +def test_expand_directories_only_lists_the_files_of_an_imas_backend(tmp_path: Path): + directory = tmp_path / "run" + directory.mkdir() + for name in ("master.h5", "equilibrium.h5", "notes.txt"): + (directory / name).write_text(name) + file = _file_data(f"imas:hdf5?path={directory}", file_type="IMAS") + + expanded = _expand_directories([file], {"data": str(tmp_path)}) + + assert [f.uri for f in expanded] == [ + "data:run/equilibrium.h5", + "data:run/master.h5", + ] + + +def test_expand_directories_reports_an_unusable_imas_uri(tmp_path: Path): + file = _file_data(f"imas:hdf5?path={tmp_path / 'missing'}", file_type="IMAS") + + with pytest.raises(APIError, match="Failed to list IMAS files"): + _expand_directories([file], {"data": str(tmp_path)}) + + +def test_expand_directories_without_configured_partitions(tmp_path: Path): + source = tmp_path / "x.txt" + source.write_text("contents") + + with pytest.raises(APIError, match="configured partitions: none"): + _expand_directories([_file_data(f"file:{source}")], {}) diff --git a/tests/remote/api/v1.3/test_resumable_upload.py b/tests/remote/api/v1.3/test_resumable_upload.py new file mode 100644 index 00000000..486c7f4f --- /dev/null +++ b/tests/remote/api/v1.3/test_resumable_upload.py @@ -0,0 +1,395 @@ +"""Tests for the resumable HTTP upload endpoint (/v1.3/upload/).""" + +import base64 +import hashlib +import uuid +from pathlib import Path +from urllib.parse import urlsplit + +import pytest +import requests +from conftest import HEADERS + +from simdb.cli import resumable_upload as ru + +INTEROP_HEADER = "Upload-Draft-Interop-Version" + + +def _digest(data): + """RFC 9530 ``sha-256`` digest structured-field value for ``data``.""" + encoded = base64.b64encode(hashlib.sha256(data).digest()).decode("ascii") + return f"sha-256=:{encoded}:" + + +@pytest.fixture +def http_partition(client, tmp_path): + """Point the ``http`` partition at a temporary directory for the test.""" + base = tmp_path / "http_staging" + base.mkdir() + client.application.simdb_config.set_option("partition.http", str(base)) + return base + + +def _patch(client, target, offset, data, complete, headers=None): + h = dict(headers or HEADERS) + h["Upload-Offset"] = str(offset) + h["Upload-Complete"] = "?1" if complete else "?0" + h["Content-Type"] = "application/partial-upload" + return client.patch(f"/v1.3/upload/{target}", data=data, headers=h) + + +def test_upload_create_append_complete(client, http_partition): + sim_hex = uuid.uuid4().hex + target = f"{sim_hex}/sub/file.txt" + + # Create the upload resource (empty body). + rv = client.post( + f"/v1.3/upload/{target}", + data=b"", + headers={**HEADERS, "Upload-Complete": "?0", "Upload-Length": "11"}, + ) + assert rv.status_code == 201 + assert rv.headers["Upload-Offset"] == "0" + assert rv.headers[INTEROP_HEADER] == "8" + assert "Location" in rv.headers + + # First chunk. + rv = _patch(client, target, 0, b"hello", complete=False) + assert rv.status_code == 204 + assert rv.headers["Upload-Offset"] == "5" + + # HEAD reports current progress. + rv = client.head(f"/v1.3/upload/{target}", headers=HEADERS) + assert rv.status_code == 204 + assert rv.headers["Upload-Offset"] == "5" + assert rv.headers["Upload-Complete"] == "?0" + + # Final chunk completes the upload. + rv = _patch(client, target, 5, b" world", complete=True) + assert rv.status_code == 200 + assert rv.headers["Upload-Offset"] == "11" + assert rv.headers["Upload-Complete"] == "?1" + + final = http_partition / sim_hex / "sub" / "file.txt" + assert final.read_bytes() == b"hello world" + assert not (final.parent / (final.name + ".partial")).exists() + + +def test_upload_offset_mismatch_returns_409(client, http_partition): + target = f"{uuid.uuid4().hex}/file.txt" + client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + _patch(client, target, 0, b"abc", complete=False) + + # Wrong offset -> 409 with the server's actual offset. + rv = _patch(client, target, 0, b"def", complete=False) + assert rv.status_code == 409 + assert rv.headers["Upload-Offset"] == "3" + + +def test_upload_head_missing_returns_404(client, http_partition): + rv = client.head(f"/v1.3/upload/{uuid.uuid4().hex}/missing.txt", headers=HEADERS) + assert rv.status_code == 404 + + +def test_upload_empty_file(client, http_partition): + sim_hex = uuid.uuid4().hex + target = f"{sim_hex}/empty.txt" + client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + rv = _patch(client, target, 0, b"", complete=True) + assert rv.status_code == 200 + assert (http_partition / sim_hex / "empty.txt").read_bytes() == b"" + + +def test_upload_delete(client, http_partition): + sim_hex = uuid.uuid4().hex + target = f"{sim_hex}/file.txt" + client.post( + f"/v1.3/upload/{target}", + data=b"data", + headers={**HEADERS, "Upload-Complete": "?1"}, + ) + assert (http_partition / sim_hex / "file.txt").exists() + + rv = client.delete(f"/v1.3/upload/{target}", headers=HEADERS) + assert rv.status_code == 204 + assert not (http_partition / sim_hex / "file.txt").exists() + + +def test_upload_path_traversal_rejected(client, http_partition): + rv = client.post( + "/v1.3/upload/..%2f..%2fescape.txt", + data=b"x", + headers={**HEADERS, "Upload-Complete": "?1"}, + ) + assert rv.status_code in (400, 403, 404) + assert not (http_partition.parent / "escape.txt").exists() + assert not Path("/tmp/escape.txt").exists() + + +class _Resp: + """Adapt a Flask test-client response to the bits resumable_upload uses.""" + + def __init__(self, rv): + self.status_code = rv.status_code + self.headers = rv.headers + self.text = rv.get_data(as_text=True) + + +def _flask_transport(client, fail_once_at=None): + """Route resumable_upload's ``requests`` calls to the Flask test client. + + @param fail_once_at: if set, raise ConnectionError the first time a PATCH is + sent at this offset, to exercise the resume path. + """ + state = {"failed": False} + + def _path(url): + return urlsplit(url).path + + def head(url, headers=None, **kwargs): + return _Resp(client.head(_path(url), headers=dict(headers or {}, **HEADERS))) + + def post(url, data=b"", headers=None, **kwargs): + return _Resp( + client.post(_path(url), data=data, headers=dict(headers or {}, **HEADERS)) + ) + + def patch(url, data=b"", headers=None, **kwargs): + offset = int((headers or {}).get("Upload-Offset", -1)) + if fail_once_at is not None and offset == fail_once_at and not state["failed"]: + state["failed"] = True + raise requests.ConnectionError("simulated network drop") + return _Resp( + client.patch(_path(url), data=data, headers=dict(headers or {}, **HEADERS)) + ) + + return head, post, patch + + +def test_resumable_upload_client_against_server(client, http_partition, monkeypatch): + head, post, patch = _flask_transport(client) + monkeypatch.setattr(ru.requests, "head", head) + monkeypatch.setattr(ru.requests, "post", post) + monkeypatch.setattr(ru.requests, "patch", patch) + + sim_hex = uuid.uuid4().hex + src = http_partition.parent / "source.bin" + payload = b"0123456789" * 100 # 1000 bytes + src.write_bytes(payload) + + ru.resumable_upload( + f"http://localhost/v1.3/upload/{sim_hex}/source.bin", src, chunk_size=256 + ) + + assert (http_partition / sim_hex / "source.bin").read_bytes() == payload + + +def test_resumable_upload_reports_progress(client, http_partition, monkeypatch): + head, post, patch = _flask_transport(client) + monkeypatch.setattr(ru.requests, "head", head) + monkeypatch.setattr(ru.requests, "post", post) + monkeypatch.setattr(ru.requests, "patch", patch) + + sim_hex = uuid.uuid4().hex + src = http_partition.parent / "progress.bin" + payload = b"y" * 1000 + src.write_bytes(payload) + + seen = [] + ru.resumable_upload( + f"http://localhost/v1.3/upload/{sim_hex}/progress.bin", + src, + chunk_size=256, + progress=seen.append, + ) + + # Progress is monotonic non-decreasing and reaches the full file size. + assert seen == sorted(seen) + assert seen[-1] == len(payload) + + +def test_resumable_upload_client_resumes_after_failure( + client, http_partition, monkeypatch +): + # Inject a connection drop at offset 256; the client should HEAD to recover + # the server offset and continue rather than restart. + head, post, patch = _flask_transport(client, fail_once_at=256) + monkeypatch.setattr(ru.requests, "head", head) + monkeypatch.setattr(ru.requests, "post", post) + monkeypatch.setattr(ru.requests, "patch", patch) + + sim_hex = uuid.uuid4().hex + src = http_partition.parent / "resume.bin" + payload = bytes(range(256)) * 4 # 1024 bytes + src.write_bytes(payload) + + ru.resumable_upload( + f"http://localhost/v1.3/upload/{sim_hex}/resume.bin", src, chunk_size=256 + ) + + assert (http_partition / sim_hex / "resume.bin").read_bytes() == payload + + +@pytest.fixture +def small_append_limit(client): + """Advertise (and enforce) a tiny max-append-size for the duration of a test.""" + cfg = client.application.simdb_config + cfg.set_option("server.max_append_size", "256") + yield 256 + cfg.set_option("server.max_append_size", str(8 * 1024 * 1024)) + + +def test_upload_advertises_and_enforces_append_limit( + client, http_partition, small_append_limit +): + target = f"{uuid.uuid4().hex}/file.bin" + rv = client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + assert rv.status_code == 201 + assert rv.headers["Upload-Limit"] == "max-append-size=256" + + # A PATCH body larger than the advertised limit is rejected. + rv = _patch(client, target, 0, b"x" * 300, complete=False) + assert rv.status_code == 413 + + +def test_client_respects_server_append_limit( + client, http_partition, small_append_limit, monkeypatch +): + head, post, patch = _flask_transport(client) + monkeypatch.setattr(ru.requests, "head", head) + monkeypatch.setattr(ru.requests, "post", post) + monkeypatch.setattr(ru.requests, "patch", patch) + + sim_hex = uuid.uuid4().hex + src = http_partition.parent / "big.bin" + payload = b"z" * 1000 # larger than the 256-byte append limit + src.write_bytes(payload) + + # Request a chunk size far larger than the server allows; the client must + # clamp to the advertised max-append-size, so the upload still succeeds. + ru.resumable_upload( + f"http://localhost/v1.3/upload/{sim_hex}/big.bin", src, chunk_size=1_000_000 + ) + + assert (http_partition / sim_hex / "big.bin").read_bytes() == payload + + +def test_resumable_upload_completes_multi_chunk(client, http_partition, monkeypatch): + head, post, patch = _flask_transport(client) + monkeypatch.setattr(ru.requests, "head", head) + monkeypatch.setattr(ru.requests, "post", post) + monkeypatch.setattr(ru.requests, "patch", patch) + + sim_hex = uuid.uuid4().hex + src = http_partition.parent / "reuse.bin" + payload = b"0123456789" * 100 + src.write_bytes(payload) + + # The file is uploaded across several chunks and assembled on the server. + ru.resumable_upload( + f"http://localhost/v1.3/upload/{sim_hex}/reuse.bin", + src, + chunk_size=256, + ) + assert (http_partition / sim_hex / "reuse.bin").read_bytes() == payload + + +def test_patch_content_digest_match_accepted(client, http_partition): + target = f"{uuid.uuid4().hex}/file.txt" + client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + rv = _patch( + client, + target, + 0, + b"hello", + complete=False, + headers={**HEADERS, "Content-Digest": _digest(b"hello")}, + ) + assert rv.status_code == 204 + assert rv.headers["Upload-Offset"] == "5" + + +def test_patch_content_digest_mismatch_rejected(client, http_partition): + target = f"{uuid.uuid4().hex}/file.txt" + client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + # Digest of different bytes than the body -> 400 and nothing appended. + rv = _patch( + client, + target, + 0, + b"hello", + complete=False, + headers={**HEADERS, "Content-Digest": _digest(b"goodbye")}, + ) + assert rv.status_code == 400 + assert rv.headers["Upload-Offset"] == "0" + assert not (http_partition / target).exists() + # The partial exists (created empty by POST) but the rejected body was not + # appended. + assert (http_partition / (target + ".partial")).read_bytes() == b"" + + +def test_multi_chunk_upload_finalizes(client, http_partition): + sim_hex = uuid.uuid4().hex + target = f"{sim_hex}/file.txt" + client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + _patch( + client, + target, + 0, + b"hello", + complete=False, + headers={**HEADERS, "Content-Digest": _digest(b"hello")}, + ) + rv = _patch( + client, + target, + 5, + b" world", + complete=True, + headers={**HEADERS, "Content-Digest": _digest(b" world")}, + ) + assert rv.status_code == 200 + assert (http_partition / sim_hex / "file.txt").read_bytes() == b"hello world" + + +def test_post_content_digest_mismatch_rejected(client, http_partition): + sim_hex = uuid.uuid4().hex + target = f"{sim_hex}/file.txt" + rv = client.post( + f"/v1.3/upload/{target}", + data=b"hello", + headers={**HEADERS, "Upload-Complete": "?0", "Content-Digest": _digest(b"x")}, + ) + assert rv.status_code == 400 + assert not (http_partition / sim_hex / "file.txt.partial").exists() + + +def test_unknown_digest_algorithm_ignored(client, http_partition): + # A digest using an algorithm the server cannot recompute is ignored rather + # than rejected, so the upload still succeeds. + target = f"{uuid.uuid4().hex}/file.txt" + client.post( + f"/v1.3/upload/{target}", data=b"", headers={**HEADERS, "Upload-Complete": "?0"} + ) + rv = _patch( + client, + target, + 0, + b"hello", + complete=False, + headers={**HEADERS, "Content-Digest": "unixsum=:0061:"}, + ) + assert rv.status_code == 204 diff --git a/tests/remote/api/v1.3/test_simulations3.py b/tests/remote/api/v1.3/test_simulations3.py index 9fb6d9b6..ab3b63bc 100644 --- a/tests/remote/api/v1.3/test_simulations3.py +++ b/tests/remote/api/v1.3/test_simulations3.py @@ -9,6 +9,7 @@ generate_simulation_data, ) +from simdb.checksum import hash_file from simdb.cli.manifest import Manifest from simdb.config import Config from simdb.database.models import Simulation @@ -19,7 +20,6 @@ ) from simdb.workers import tasks as simdb_tasks from simdb.workers.celery import celery_app -from simdb.workers.tasks import _calculate_checksum @pytest.fixture(autouse=True) @@ -82,7 +82,7 @@ def generate_simulation_file(path) -> FileData: file_path = path / "partition/file.txt" file_path.parent.mkdir(exist_ok=True) file_path.write_text("test data") - checksum = _calculate_checksum(file_path) + checksum = hash_file(file_path) return FileData( type="FILE", uri="data:///file.txt", diff --git a/tests/validation/test_validator.py b/tests/validation/test_validator.py new file mode 100644 index 00000000..bcef6a40 --- /dev/null +++ b/tests/validation/test_validator.py @@ -0,0 +1,56 @@ +import numpy as np + +from simdb.validation.validator import CustomValidator + + +def test_custom_validator_min_value_max_value(): + schema = { + "field1": { + "type": "numpy", + "coerce": "numpy", + "min_value": 0.0, + "max_value": 10.0, + } + } + validator = CustomValidator(schema) + + # Test valid numpy array + assert validator.validate({"field1": np.array([1.0, 5.0, 9.0])}) + + # Test valid dictionary representing a range + assert validator.validate({"field1": {"min": 1.0, "max": 9.0}}) + + # Test numpy array out of bounds (too low) + assert not validator.validate({"field1": np.array([-1.0, 5.0, 9.0])}) + + # Test numpy array out of bounds (too high) + assert not validator.validate({"field1": np.array([1.0, 5.0, 11.0])}) + + # Test dictionary range out of bounds (min too low) + assert not validator.validate({"field1": {"min": -1.0, "max": 9.0}}) + + # Test dictionary range out of bounds (max too high) + assert not validator.validate({"field1": {"min": 1.0, "max": 11.0}}) + + +def test_custom_validator_comparisons(): + schema = { + "field_ge": {"type": "numpy", "coerce": "numpy", "ge": 0.0}, + "field_le": {"type": "numpy", "coerce": "numpy", "le": 10.0}, + } + validator = CustomValidator(schema) + + # Test valid dictionary representing a range + assert validator.validate( + {"field_ge": {"min": 0.0, "max": 5.0}, "field_le": {"min": 1.0, "max": 10.0}} + ) + + # Test invalid range for ge (min is -1, which is not >= 0) + assert not validator.validate( + {"field_ge": {"min": -1.0, "max": 5.0}, "field_le": {"min": 1.0, "max": 10.0}} + ) + + # Test invalid range for le (max is 11, which is not <= 10) + assert not validator.validate( + {"field_ge": {"min": 0.0, "max": 5.0}, "field_le": {"min": 1.0, "max": 11.0}} + ) diff --git a/tests/workers/test_tasks.py b/tests/workers/test_tasks.py index 6d411be4..22de3a33 100644 --- a/tests/workers/test_tasks.py +++ b/tests/workers/test_tasks.py @@ -1,16 +1,18 @@ +import hashlib from datetime import datetime, timezone from unittest import mock from uuid import uuid1 import pytest +from simdb.checksum import hash_file from simdb.config import Config from simdb.enums import IngestionStatus from simdb.imas.utils import SimDBUrl from simdb.remote.models import FileData from simdb.workers import tasks as simdb_tasks from simdb.workers.tasks import ( - _calculate_checksum, + _checksum_matches, _copy_files, _create_file_from_data, _get_imas_identifier_path, @@ -18,6 +20,7 @@ _notify_watchers, _resolve_paths, _resolve_uri_to_path, + cleanup_http_staging_task, copy_files_task, ) @@ -149,6 +152,15 @@ def test_create_file_from_data_raises_on_checksum_mismatch(tmp_path): _create_file_from_data(file_data, config, data_file) +def test_checksum_matches_uses_sha1(tmp_path): + data_file = tmp_path / "testfile.txt" + content = b"content" + data_file.write_bytes(content) + + assert _checksum_matches(data_file, hashlib.sha1(content).hexdigest()) + assert not _checksum_matches(data_file, hashlib.sha1(b"other").hexdigest()) + + @pytest.fixture def task_environment(tmp_path): """Set up Config, mocked DB, and directory layout for copy_files_task tests.""" @@ -162,7 +174,7 @@ def task_environment(tmp_path): config.set_option("database.file", str(tmp_path / "test.db")) config.set_option("server.upload_folder", str(upload_dir)) config.set_option("partition.data", str(partition_dir)) - config.load = mock.MagicMock() + config.load = mock.MagicMock() # ty: ignore[invalid-assignment] simulation_uuid = uuid1() mock_simulation = mock.MagicMock(uuid=simulation_uuid, inputs=[], outputs=[]) @@ -188,9 +200,7 @@ def test_copy_files_task_copies_inputs_and_marks_copied(task_environment): source_file.write_text("test content") input_files = [ - _make_file_data( - f"data:/{source_file.name}", checksum=_calculate_checksum(source_file) - ) + _make_file_data(f"data:/{source_file.name}", checksum=hash_file(source_file)) ] copy_files_task(env["simulation_uuid"], input_files, []) @@ -226,3 +236,91 @@ def test_notify_watchers_noop_without_watchers(): _notify_watchers(simulation, "subject", "body") delay.assert_not_called() + + +def test_copy_files_task_http_keeps_imas_folder_flattens_sibling(task_environment): + """HTTP-staged files are copied like local push: a shared root is stripped, + so an IMAS directory keeps its folder while a sibling file stays flat.""" + env = task_environment + sim_hex = env["simulation_uuid"].hex + + http_partition = env["partition_dir"].parent / "http_staging" + env["config"].set_option("partition.http", str(http_partition)) + + # Stage as the client would: /data/subdir/{test_hdf5/*, test.nc} + staged = http_partition / sim_hex / "data" / "subdir" + (staged / "test_hdf5").mkdir(parents=True) + master = staged / "test_hdf5" / "master.h5" + extra = staged / "test_hdf5" / "0001.h5" + nc = staged / "test.nc" + master.write_text("m") + extra.write_text("d") + nc.write_text("n") + + output_files = [ + _make_file_data( + f"http://{sim_hex}/data/subdir/test_hdf5/master.h5", + checksum=hash_file(master), + ), + _make_file_data( + f"http://{sim_hex}/data/subdir/test_hdf5/0001.h5", + checksum=hash_file(extra), + ), + _make_file_data( + f"http://{sim_hex}/data/subdir/test.nc", checksum=hash_file(nc) + ), + ] + + copy_files_task(env["simulation_uuid"], [], output_files) + + dest = env["upload_dir"] / sim_hex + # The IMAS hdf5 directory keeps its folder... + assert (dest / "test_hdf5" / "master.h5").read_text() == "m" + assert (dest / "test_hdf5" / "0001.h5").read_text() == "d" + # ...while the standalone netcdf file is not given a spurious parent folder. + assert (dest / "test.nc").read_text() == "n" + assert env["simulation"].ingestion_status == IngestionStatus.COPIED + + +def test_resolve_uri_to_path_folds_http_host_into_path(tmp_path): + """http:// URIs put the sim-uuid in the authority; it must be reconstructed.""" + config = Config() + partition_path = tmp_path / "http_staging" + partition_path.mkdir() + config.set_option("partition.http", str(partition_path)) + + uri = SimDBUrl("http://deadbeef/subdir/file.txt") + result = _resolve_uri_to_path(uri, config) + + assert result == partition_path / "deadbeef" / "subdir" / "file.txt" + + +def test_cleanup_http_staging_task_removes_simulation_dir(tmp_path): + partition_path = tmp_path / "http_staging" + sim_uuid = uuid1() + staging = partition_path / sim_uuid.hex + staging.mkdir(parents=True) + (staging / "file.txt").write_text("data") + # A sibling simulation's data must be left untouched. + other = partition_path / "other" + other.mkdir() + (other / "keep.txt").write_text("keep") + + config = Config() + config.set_option("partition.http", str(partition_path)) + config.load = mock.MagicMock() # ty: ignore[invalid-assignment] + + with mock.patch("simdb.workers.tasks.Config", return_value=config): + cleanup_http_staging_task(sim_uuid) + + assert not staging.exists() + assert (other / "keep.txt").exists() + + +def test_cleanup_http_staging_task_without_partition_is_noop(tmp_path): + config = Config() + config.load = mock.MagicMock() # ty: ignore[invalid-assignment] + + with mock.patch("simdb.workers.tasks.Config", return_value=config): + # Should not raise even though partition.http is unset. + cleanup_http_staging_task(uuid1()) diff --git a/uv.lock b/uv.lock index 4130318a..581d0cf7 100644 --- a/uv.lock +++ b/uv.lock @@ -1266,7 +1266,7 @@ requires-dist = [ { name = "imas-simdb", extras = ["imas-validator", "postgres", "server"], marker = "extra == 'all'" }, { name = "imas-validator", marker = "extra == 'imas-validator'", specifier = ">=1.0.0" }, { name = "myst-parser", marker = "extra == 'build-docs'", specifier = ">=0.18.0" }, - { name = "netcdf4", specifier = ">=1.5" }, + { name = "netcdf4", specifier = ">=1.7.2" }, { name = "numpy", specifier = ">=1.14" }, { name = "plotext", specifier = "==5.3.2" }, { name = "psycopg2-binary", marker = "extra == 'postgres'", specifier = ">=2.8.0" },