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/cli.md b/docs/cli.md index 37875bdf..62f3c64f 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -14,7 +14,6 @@ Options: Commands: alias Query remote and local aliases. config Query/update application configuration. - database Manage local simulation database. manifest Create/check manifest file. provenance Create the PROVENANCE_FILE from the current system. remote Interact with the remote SimDB service. @@ -142,7 +141,6 @@ Options: --help Show this message and exit. ``` - ## Manifest @@ -421,30 +419,38 @@ Usage: simdb remote [NAME] query [OPTIONS] [CONSTRAINTS]... NAME=[mod]VALUE Where `[mod]` is an optional query modifier. Available query modifiers are: - eq: - This checks for equality (this is the same behaviour as not providing any modifier). + eq: - This checks for equality (this is the same behaviour as not providing + any modifier). in: - This searches inside the value instead of looking for exact matches. gt: - This checks for values greater than the given quantity. agt: - This checks for any array elements are greater than the given quantity. ge: - This checks for values greater than or equal to the given quantity. - age: - This checks for any array elements are greater than or equal to the given quantity. + age: - This checks for any array elements are greater than or equal to the given + quantity. lt: - This checks for values less than the given quantity. - alt: - This checks for any array elements are less than the given quantity. + alt: - This checks for any array elements are less than the given quantity. le: - This checks for values less than or equal to the given quantity. - ale: - This checks for any array elements are less than or equal to the given quantity. + ale: - This checks for any array elements are less than or equal to the given + quantity. Modifier examples: alias=eq:foo performs exact match - summary.code.name=in:foo matches all names containing foo - summary.heating_current_drive.power_additional.value=agt:0 matches all simulations where any array element - of summary.heating_current_drive.power_additional.value is greater than 0 + summary.code.name=in:foo matches all names + containing foo + summary.heating_current_drive.power_additional.value=agt:0 matches all + simulations where any array element of + summary.heating_current_drive.power_additional.value is greater than 0 - Any string comparisons are done in a case-insensitive manner. If multiple constraints are provided then simulations - are returned that match all given constraints. + Any string comparisons are done in a case-insensitive manner. If multiple + constraints are provided then simulations are returned that match all given + constraints. Examples: - sim remote query workflow.name=in:test finds all simulations where workflow.name contains test - (case-insensitive) - sim remote query pulse=gt:1000 run=0 finds all simulations where pulse is > 1000 and run = 0 + sim remote query workflow.name=in:test finds all simulations where + workflow.name contains test + (case-insensitive) + sim remote query pulse=gt:1000 run=0 finds all simulations where pulse + is > 1000 and run = 0 Options: -m, --meta-data TEXT Additional meta-data field to print. @@ -603,20 +609,44 @@ Options: --help Show this message and exit. Commands: - delete Delete the ingested simulation with given SIM_ID (UUID or... - info Print information on the simulation with given SIM_ID (UUID... - ingest Ingest a MANIFEST_FILE. - list List ingested simulations. - modify Modify the ingested simulation. - pull Pull the simulation with the given SIM_ID (UUID or alias)... - push Push the simulation with the given SIM_ID (UUID or alias) to... - query Perform a metadata query to find matching local simulations. - validate Validate the ingested simulation with given SIM_ID (UUID or... + data Fetch IDS field data for simulation SIM_ID (UUID or alias)... + delete Delete the ingested simulation with given SIM_ID (UUID or... + info Print information on the simulation with given SIM_ID (UUID... + ingest Ingest a MANIFEST_FILE. + list List ingested simulations. + modify Modify the ingested simulation. + pull Pull the simulation with the given SIM_ID (UUID or alias)... + push Push the simulation with the given SIM_ID (UUID or alias)... + push_http Push the simulation with the given SIM_ID to the REMOTE... + push_local Push the simulation with the given SIM_ID (UUID or alias)... + query Perform a metadata query to find matching local simulations. + validate Validate the ingested simulation with given SIM_ID (UUID or... ``` ```text -Usage: simdb simulation delete [OPTIONS] SIM_ID +Usage: simdb simulation data [OPTIONS] [REMOTE] SIM_ID IDS_PATH + + Fetch IDS field data for simulation SIM_ID (UUID or alias) from REMOTE. + + IDS_PATH format: + ids_name[:]/path/to/field + + Examples: + simdb sim data iter 4dd781b... profiles_1d[0]/grid/rho_tor_norm + simdb sim data 4dd781b... equilibrium:0/time_slice[0]/profiles_1d/psi + +Options: + --username TEXT Username used to authenticate with the remote. + --password TEXT Password used to authenticate with the remote. + --dd-version TEXT Convert IDS data to the requested Data Dictionary + version, e.g. 4.1.1. + --help Show this message and exit. +``` + + +```text +Usage: simdb simulation delete [OPTIONS] [SIM_ID] Delete the ingested simulation with given SIM_ID (UUID or alias). @@ -703,6 +733,41 @@ Options: ``` +```text +Usage: simdb simulation push_http [OPTIONS] [REMOTE] SIM_ID + + Push the simulation with the given SIM_ID to the REMOTE over resumable HTTP. + + Unlike push_local, this does not require a filesystem shared with the + server: the file bytes are uploaded over HTTP using a resumable protocol and + staged into the server's 'http' partition. An interrupted push can be + resumed by re-running the command. + +Options: + --username TEXT Username used to authenticate with the remote. + --password TEXT Password used to authenticate with the remote. + --replaces TEXT SIM_ID of simulation to deprecate and replace. + --add-watcher Add the current user as a watcher of the simulation. + --help Show this message and exit. +``` + + +```text +Usage: simdb simulation push_local [OPTIONS] [REMOTE] SIM_ID + + Push the simulation with the given SIM_ID (UUID or alias) to the REMOTE. + +Options: + --username TEXT Username used to authenticate with the remote. + --password TEXT Password used to authenticate with the remote. + --replaces TEXT SIM_ID of simulation to deprecate and replace. + --add-watcher Add the current user as a watcher of the simulation. + --timeout FLOAT Maximum number of seconds to wait for ingestion to + complete. [default: 600.0] + --help Show this message and exit. +``` + + ```text Usage: simdb simulation query [OPTIONS] [CONSTRAINTS]... @@ -712,7 +777,8 @@ Usage: simdb simulation query [OPTIONS] [CONSTRAINTS]... NAME=[mod]VALUE Where `[mod]` is an optional query modifier. Available query modifiers are: - eq: - This checks for equality (this is the same behaviour as not providing any modifier). + eq: - This checks for equality (this is the same behaviour as not providing any + modifier). ne: - This checks for value that do not equal. in: - This searches inside the value instead of looking for exact matches. ni: - This searches inside the value for elements that do not match. @@ -722,22 +788,26 @@ Usage: simdb simulation query [OPTIONS] [CONSTRAINTS]... le: - This checks for values less than or equal to the given quantity. For the following modifiers, VALUE should not be provided. exist: - This - returns simulations where metadata with NAME exists, regardless of the - value. + returns simulations where metadata with NAME exists, regardless + of the value. Modifier examples: responsible_name=foo performs exact match responsible_name=in:foo matches all names containing foo pulse=gt:1000 matches all pulses > 1000 - sequence=exist: matches all simulations that have "sequence" metadata values + sequence=exist: matches all simulations that have "sequence" + metadata values - Any string comparisons are done in a case-insensitive manner. If multiple constraints are provided then simulations - are returned that match all given constraints. + Any string comparisons are done in a case-insensitive manner. If multiple + constraints are provided then simulations are returned that match all given + constraints. Examples: - sim simulation query workflow.name=in:test finds all simulations where workflow.name contains test + sim simulation query workflow.name=in:test finds all simulations where + workflow.name contains test (case-insensitive) - sim simulation query pulse=gt:1000 run=0 finds all simulations where pulse is > 1000 and run = 0 + sim simulation query pulse=gt:1000 run=0 finds all simulations where + pulse is > 1000 and run = 0 Options: -m, --meta-data TEXT Additional meta-data field to print. diff --git a/docs/cli.md.in b/docs/cli.md.in index c1b8371b..adf14fa3 100644 --- a/docs/cli.md.in +++ b/docs/cli.md.in @@ -10,10 +10,6 @@ {{ config }} -## Database - -{{ database }} - ## Manifest {{ manifest }} diff --git a/docs/user_guide.md b/docs/user_guide.md index 15b927b0..6dcfb12e 100644 --- a/docs/user_guide.md +++ b/docs/user_guide.md @@ -297,6 +297,83 @@ simdb simulation push This will upload all the metadata associated with your simulation to the remote server as well as taking copies of all input and output data specified. For non-IMAS data the `file` URIs will be used to locate the files to transfer, whereas for `imas` URIs SimDB will discover which files need to be transferred based on the IMAS backend specified in the URI. The files are copied to the server using an HTTP data transfer. +### Pushing Local Simulations (Optimized for Shared File Systems) + +If your local environment and the remote SimDB server share a common local file system (for example, on the ITER network where both you and the server can access same physical file paths directly), uploading large datasets over HTTP can be slow and redundant. + +In this scenario, you should use the `push_local` command: + +```bash +simdb simulation push_local +``` + +Unlike the standard `push` command, `push_local` only sends the simulation metadata and the storage file paths to the server. The remote server then: +1. Validates the simulation metadata against active schemas. +2. Queues the file copy operation in the background using an asynchronous Celery task queue. +3. Automatically completes the ingestion once background file copying finishes. + +The CLI command will block and print real-time updates while waiting for the background ingestion to complete: +```text +Waiting for ingestion to complete... queued -> copy_files -> completed +Successfully pushed simulation +``` + +#### Configuring Partitions (Shared Storage Mappings) + +To enable `push_local` to map and resolve files correctly between your local environment and the remote server, you must configure **partitions** in your `simdb.cfg` configuration file. + +Partitions define mappings between short, logical names (such as `data` or `work`) and their absolute paths on your local file system. + +##### 1. Defining Partitions on the Client +Add a `[partition]` section to your `~/.config/simdb/simdb.cfg` file, specifying the directory paths for each partition: + +```ini +[partition] +data = /home/user/my_simdb_data +work = /work/imas/shared +sdcc = / +``` + +*Note on `sdcc` partition mapping:* +In environments like the ITER network, files are often located under absolute paths like `/sdcc/projects/...`. Mapping the `sdcc` partition to the system root `/` ensures that any path beginning with `/sdcc` is correctly matched and converted to a partition-relative URI (e.g., `/sdcc/projects/my_run` becomes `sdcc:sdcc/projects/my_run`). When several partitions contain a file, the partition with the most specific (deepest) path wins, so such a catch-all mapping never shadows the other partitions. + +##### 2. How Partitions are Resolved +When you run `push_local`: +* SimDB scans the manifest's input/output files and checks if any path falls under one of your defined local partitions. +* If a match is found (e.g., `/home/user/my_simdb_data/scenarios/run1.txt` is inside `/home/user/my_simdb_data`), SimDB converts the file URI into a partition-relative scheme: `data:scenarios/run1.txt`. +* The remote server receives this logical URI. As long as the server also has the `data` partition configured (even if mounted at a different absolute path like `/mnt/shared/partition`), it resolves the URI to `/mnt/shared/partition/scenarios/run1.txt` and completes ingestion. + +This mapping mechanism allows clients and servers to share data over a network or cluster filesystem even if they mount it at different absolute paths. + +### Pushing Simulations over Resumable HTTP + +When your local environment and the remote server do **not** share a filesystem, but you still want the background ingestion workflow of `push_local` (rather than the single, non-resumable transfer of `push`), use the `push_http` command: + +```bash +simdb simulation push_http +``` + +`push_http` uploads each file's bytes to the server over HTTP using the IETF "Resumable Uploads for HTTP" protocol (draft-ietf-httpbis-resumable-upload-11). The uploaded files are staged into a server-side partition named `http`, and from then on the flow is identical to `push_local`: SimDB sends the metadata, the server queues a background copy into its upload folder, and the CLI blocks while printing ingestion status updates. + +Because the protocol is resumable, an interrupted upload (lost connection, Ctrl-C) does not have to start over. Re-running `push_http` for the same simulation asks the server how many bytes it already received for each file and continues from that offset. + +The size of each upload chunk is governed by the server: it advertises a maximum append size via 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 on the server with the `server.max_append_size` config option (for example, to fit within a reverse proxy's request body limit). + +Each chunk is integrity-checked using the RFC 9530 digest field. The client sends a `Content-Digest` (SHA-256) with every chunk so the server can verify that 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. + +#### Configuring the `http` Partition on the Server + +Unlike `push_local`, `push_http` needs **no client-side partition configuration** - the file paths are taken directly from the local simulation. The server, however, must define where uploaded bytes are staged by configuring a partition named `http` in its `simdb.cfg`: + +```ini +[partition] +http = /var/lib/simdb/http-staging +``` + +When a file is uploaded to `/`, the server writes it to `//` and references it with an `http:///` 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 from the `http` partition. + +Subfolder structure is handled exactly as it is for `push_local`: files keep their partition-relative layout, so multi-file IMAS datasets (HDF5, ASCII, 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. + ## Pulling simulations from a remote The mirror to pushing simulations is the `pull` command. This command will pull the simulation metadata from the SimDB remote to your local SimDB database and download the simulation data into a directory of your choosing. Once you have pulled a simulation it will appear in any local SimDB queries you perform. The command looks as follows: diff --git a/pyproject.toml b/pyproject.toml index 0ff24106..44a29a20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,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/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..7aab39dc 100644 --- a/src/simdb/cli/commands/simulation.py +++ b/src/simdb/cli/commands/simulation.py @@ -1,5 +1,6 @@ import contextlib import sys +import time import urllib.parse from itertools import chain from pathlib import Path @@ -14,6 +15,7 @@ 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 @@ -212,6 +214,192 @@ def parse_args(self, ctx, args): return NRequiredArgs +@simulation.command("push_local", cls=n_required_args_adaptor(1)) +@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.", +) +@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, + username: Optional[str], + 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.""" + + api = RemoteAPI(remote, username, password, config) + 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 + + api.push_local_simulation(simulation, add_watcher=add_watcher) + + max_consecutive_failures = 5 + + click.echo("Waiting for ingestion to complete...", nl=False) + last_status = None + consecutive_failures = 0 + deadline = time.monotonic() + timeout + while True: + try: + status = api.get_ingestion_status(simulation.uuid.hex) + 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(1) + continue + consecutive_failures = 0 + + 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(): + break + + 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(1) + + click.echo() + if status == IngestionStatus.COMPLETED: + click.echo(f"Successfully pushed simulation {simulation.uuid}") + else: + raise click.ClickException( + f"Simulation ingestion failed with status: {status.value}" + ) + + +@simulation.command("push_http", cls=n_required_args_adaptor(1)) +@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_http( + 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 to the REMOTE over resumable HTTP. + + Unlike push_local, this does not require a filesystem shared with the server: + the file bytes are uploaded over HTTP using a resumable protocol and staged + into the server's 'http' partition. An interrupted push can be resumed by + re-running the command. + """ + + api = RemoteAPI(remote, username, password, config) + 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 + + api.push_http_simulation(simulation, add_watcher=add_watcher) + + click.echo("Waiting for ingestion to complete...", nl=False) + last_status = None + while True: + try: + status = api.get_ingestion_status(simulation.uuid.hex) + except Exception as err: + click.echo() + raise click.ClickException( + f"Failed to check ingestion status: {err}" + ) 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(): + break + + time.sleep(1) + + click.echo() + if status == IngestionStatus.COMPLETED: + click.echo(f"Successfully pushed simulation {simulation.uuid}") + else: + raise click.ClickException( + f"Simulation ingestion failed with status: {status.value}" + ) + + @simulation.command("push", cls=n_required_args_adaptor(1)) @pass_config @click.argument("remote", required=False) diff --git a/src/simdb/cli/remote_api.py b/src/simdb/cli/remote_api.py index 99f87363..abb1e290 100644 --- a/src/simdb/cli/remote_api.py +++ b/src/simdb/cli/remote_api.py @@ -25,19 +25,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 +248,197 @@ 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] +) -> FileData: + partition, partition_path = _find_partition_for_file(source, partitions) + new_uri = SimDBUrl.build(scheme=partition, path=partition_path.as_posix()) + return FileData( + type=file.type, + uri=new_uri.encoded_string(), + checksum=hash_file(source), + datetime=file.datetime, + usage=file.usage, + purpose=file.purpose, + sensitivity=file.sensitivity, + access=file.access, + embargo=file.embargo, + ) + + +def _expand_directories(files: Iterable[FileData], partitions: Dict[str, str]): + new_file_list = [] + for file in files: + file_uri = SimDBUrl(file.uri) + if file_uri.path is None: + raise APIError(f"File URI has no path: {file.uri}") + 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 APIError(f"IMAS URI has no path set: {file.uri}") + file_path = Path(path) + + if file_path.is_dir(): + for sub_file in file_path.iterdir(): + if sub_file.is_dir(): + raise APIError( + f"Nested directory found in {file_path}: {sub_file.name}" + ) + new_file_list.append( + _file_data_for_partition(file, sub_file, partitions) + ) + else: + new_file_list.append(_file_data_for_partition(file, file_path, partitions)) + 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. @@ -856,6 +1060,105 @@ def _send_chunk( ] self.post("files", data={}, files=files) + @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_http_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")) + + @versioned_method("v1.3") + @try_request + def get_ingestion_status(self, sim_id: str) -> str: + res = self.get(f"simulation/status/{sim_id}") + return IngestionStatus(res.json()["status"]) + @versioned_method("v1.2", "v1.3") @try_request def push_simulation( @@ -1058,7 +1361,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") @@ -1067,8 +1370,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) @@ -1084,7 +1387,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..368ad535 --- /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 7b9ff234..8d8538fa 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/__init__.py b/src/simdb/remote/apis/__init__.py index b75a9680..2f8a7803 100644 --- a/src/simdb/remote/apis/__init__.py +++ b/src/simdb/remote/apis/__init__.py @@ -36,6 +36,26 @@ def register(api, version, namespaces): version_str = version.replace(".", "_") blueprint = Blueprint(f"api_{version_str}", f"{__name__}.{version_str}") blueprints[version] = blueprint + + def index(): + return jsonify( + { + "api": "simdb", + "api_version": api.version, + "server_version": __version__, + "endpoints": [ + request.url + "simulations", + request.url + "files", + request.url + "validation_schema", + request.url + "metadata", + request.url + "upload_options", + ], + "documentation": request.url + "docs", + } + ) + + api.render_root = index + api.init_app(blueprint) for namespace in namespaces: @@ -78,21 +98,7 @@ def handle_authentication_error(err: Exception): class Index(Resource): @api.doc(security=[]) def get(self): - return jsonify( - { - "api": "simdb", - "api_version": api.version, - "server_version": __version__, - "endpoints": [ - request.url + "simulations", - request.url + "files", - request.url + "validation_schema", - request.url + "metadata", - request.url + "upload_options", - ], - "documentation": request.url + "docs", - } - ) + return index() @api.route("/token") class Token(Resource): 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_push_http.py b/tests/cli/test_push_http.py new file mode 100644 index 00000000..69c3bee1 --- /dev/null +++ b/tests/cli/test_push_http.py @@ -0,0 +1,181 @@ +"""Tests for the HTTP push client helpers and CLI command.""" + +import hashlib +import uuid +from datetime import datetime, timezone +from unittest import mock + +from click.testing import CliRunner +from utils import config_test_file + +from simdb.cli.remote_api import _compute_checksums, _expand_directories_http +from simdb.cli.simdb import cli +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('/')}" + + +def test_push_http_command_pushes_and_reports_success(tmp_path): + runner = CliRunner() + config_file = config_test_file() + + fake_api = mock.MagicMock() + fake_api.get_validation_schemas.return_value = [] + fake_api.get_ingestion_status.return_value = "completed" + + sim = mock.MagicMock() + sim.uuid = uuid.uuid4() + fake_db = mock.MagicMock() + fake_db.get_simulation.return_value = sim + + with mock.patch( + "simdb.cli.commands.simulation.RemoteAPI", return_value=fake_api + ), mock.patch("simdb.cli.commands.simulation.get_local_db", return_value=fake_db): + result = runner.invoke( + cli, + [f"--config-file={config_file}", "simulation", "push_http", "iter", "sim1"], + ) + + assert result.exit_code == 0, result.output + fake_api.push_http_simulation.assert_called_once_with(sim, add_watcher=False) + assert "Successfully pushed simulation" in result.output + + +def test_push_http_command_passes_add_watcher(tmp_path): + runner = CliRunner() + config_file = config_test_file() + + fake_api = mock.MagicMock() + fake_api.get_validation_schemas.return_value = [] + fake_api.get_ingestion_status.return_value = "completed" + + sim = mock.MagicMock() + sim.uuid = uuid.uuid4() + fake_db = mock.MagicMock() + fake_db.get_simulation.return_value = sim + + with mock.patch( + "simdb.cli.commands.simulation.RemoteAPI", return_value=fake_api + ), mock.patch("simdb.cli.commands.simulation.get_local_db", return_value=fake_db): + result = runner.invoke( + cli, + [ + f"--config-file={config_file}", + "simulation", + "push_http", + "iter", + "sim1", + "--add-watcher", + ], + ) + + assert result.exit_code == 0, result.output + fake_api.push_http_simulation.assert_called_once_with(sim, add_watcher=True) + + +def test_push_http_command_fails_on_failed_status(tmp_path): + runner = CliRunner() + config_file = config_test_file() + + fake_api = mock.MagicMock() + fake_api.get_validation_schemas.return_value = [] + fake_api.get_ingestion_status.return_value = "copy_failed" + + sim = mock.MagicMock() + sim.uuid = uuid.uuid4() + fake_db = mock.MagicMock() + fake_db.get_simulation.return_value = sim + + with mock.patch( + "simdb.cli.commands.simulation.RemoteAPI", return_value=fake_api + ), mock.patch("simdb.cli.commands.simulation.get_local_db", return_value=fake_db): + result = runner.invoke( + cli, + [f"--config-file={config_file}", "simulation", "push_http", "iter", "sim1"], + ) + + assert result.exit_code != 0 + assert "copy_failed" in result.output diff --git a/tests/remote/api/test_index.py b/tests/remote/api/test_index.py new file mode 100644 index 00000000..608fadcc --- /dev/null +++ b/tests/remote/api/test_index.py @@ -0,0 +1,16 @@ +import pytest + +from simdb.remote.apis import blueprints + + +@pytest.mark.parametrize("version", list(blueprints)) +def test_versioned_index(client, version): + """The versioned root must serve the index JSON, not flask-restx's 404 root.""" + rv = client.get(f"/{version}/") + + assert rv.status_code == 200 + assert rv.json["api"] == "simdb" + # v1 reports "1.0" for blueprint key "v1", so only compare the prefix + assert rv.json["api_version"].startswith(version.lstrip("v")) + assert "server_version" in rv.json + assert any(url.endswith("simulations") for url in rv.json["endpoints"]) 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 2216c930..532738e0 100644 --- a/uv.lock +++ b/uv.lock @@ -3146,7 +3146,7 @@ requires-dist = [ { name = "imas-validator", marker = "extra == 'imas-validator'", specifier = ">=1.0.0" }, { name = "myst-parser", marker = "extra == 'build-docs'", specifier = ">=0.18.0" }, { name = "nbsphinx", marker = "extra == 'build-docs'", specifier = ">=0.8.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" },