diff --git a/docs/generate_cli_docs.py b/docs/generate_cli_docs.py index faf5527e..b0a00933 100644 --- a/docs/generate_cli_docs.py +++ b/docs/generate_cli_docs.py @@ -39,7 +39,9 @@ def extract_sub_commands(output: str) -> list[str]: sub_commands = [] for line in output.split("\n"): if in_commands: - if line: + # Lines indented further than the command column are continuations + # of the previous command's help text. + if line and not line.startswith(" "): sub_commands.append(line.split()[0]) if line == "Commands:": in_commands = True diff --git a/docs/how-to/push-pull.md b/docs/how-to/push-pull.md index 6b555080..b4a94122 100644 --- a/docs/how-to/push-pull.md +++ b/docs/how-to/push-pull.md @@ -42,6 +42,55 @@ simdb simulation push SIM_ID --add-watcher See [watchers](../explanation/concepts.md#watchers) and the `simdb remote watcher` commands in the [CLI reference](../reference/cli.md). +## Push on a shared file system + +If your machine and the server can reach the same physical file paths (as on the +ITER network), sending large datasets over HTTP is slow and redundant. Use +`push_local` instead: + +```bash +simdb simulation push_local SIM_ID +``` + +`push_local` sends only the metadata and the storage paths. The server then + +1. validates the metadata against the active schemas, +2. queues the file copy as a background [Celery task](operate-server/run-celery-workers.md), and +3. completes the ingestion once the copy finishes. + +The command blocks and reports the ingestion state as it changes: + +```text +Waiting for ingestion to complete... QUEUED -> COPYING -> COPIED -> COMPLETED +Successfully pushed simulation UUID +``` + +### Configure partitions + +For `push_local` to resolve files on both sides, client and server must agree on +a set of *partitions*: short logical names mapped to absolute directories. Add a +`[partition]` section to your client configuration (see +[Client configuration](../reference/configuration.md#partition)): + +```ini +[partition] +data = /home/user/my_simdb_data +work = /work/imas/shared +sdcc = / +``` + +Mapping `sdcc` to the system root makes any path under `/sdcc/projects/...` +match, so `/sdcc/projects/my_run` becomes `sdcc:sdcc/projects/my_run`. When +several partitions contain a file the most specific (deepest) path wins, so a +catch-all mapping like this never shadows the others. + +When you run `push_local`, SimDB checks every input and output path against your +partitions. A path inside a partition is rewritten to a partition-relative URI — +`/home/user/my_simdb_data/scenarios/run1.txt` becomes +`data:scenarios/run1.txt`. The server resolves that URI against its own +`[partition]` configuration, so the two sides may mount the same storage at +different absolute paths. + ## Pull Pull copies a simulation's metadata into your local catalogue and downloads its diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index f2bbc65c..2a4aa04e 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -90,6 +90,17 @@ One section per configured remote server. Manage these with | --- | --- | | `file` | Path to the local SQLite catalogue. Defaults to `sim.db` in the user data directory (for example `~/.local/share/simdb/sim.db`). | +### `[partition]` + +Maps logical partition names to absolute directories on this machine. Used by +`simdb simulation push_local` to rewrite file paths into partition-relative URIs +that the server can resolve (see +[Push and pull simulations](../how-to/push-pull.md#configure-partitions)). + +| Option | Description | +| --- | --- | +| `NAME` | Directory that partition `NAME` is mounted at, for example `data = /home/user/my_simdb_data`. | + ### `[development]` | Option | Description | diff --git a/pyproject.toml b/pyproject.toml index 7b747abc..c30d64c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "distro>=1.8.0", "email-validator>=1.1", "imas-python>=2.0.1", - "netCDF4>=1.5", + "netCDF4>=1.7.2", "numpy>=1.14", "pydantic>=2.10.6", "python-dateutil>=2.6", diff --git a/scripts/test_v13_ingestion.py b/scripts/test_v13_ingestion.py index 630a1dfc..cca60518 100755 --- a/scripts/test_v13_ingestion.py +++ b/scripts/test_v13_ingestion.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 """Test script for v1.3 simulation ingestion against a running server.""" -from simdb.workers.tasks import _calculate_checksum +from simdb.checksum import calculate_checksum from pathlib import Path import base64 @@ -37,7 +37,7 @@ def generate_simulation_file(): - checksum = _calculate_checksum(Path("tmp/partition_data/subdir/test_file.txt")) + checksum = calculate_checksum(Path("tmp/partition_data/subdir/test_file.txt")) return FileData( type="FILE", uri="data:///subdir/test_file.txt", @@ -46,7 +46,7 @@ def generate_simulation_file(): ) def generate_imas_file(relative_path): - checksum = _calculate_checksum(Path(f"tmp/partition_data/subdir/{relative_path}")) + checksum = calculate_checksum(Path(f"tmp/partition_data/subdir/{relative_path}")) return FileData( type="IMAS", uri=f"data:///subdir/{relative_path}", diff --git a/src/simdb/checksum.py b/src/simdb/checksum.py index 99aef230..f85aea1f 100644 --- a/src/simdb/checksum.py +++ b/src/simdb/checksum.py @@ -4,6 +4,19 @@ from simdb.imas.utils import SimDBUrl +def calculate_checksum(path: Path) -> str: + """Generate a SHA1 checksum from the file at the given path. + + :param path: the path of the file to checksum + :return: a string containing the hex representation of the computed SHA1 checksum + """ + sha1 = hashlib.sha1() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(4096), b""): + sha1.update(chunk) + return sha1.hexdigest() + + def sha1_checksum(uri: SimDBUrl) -> str: """Generate a SHA1 checksum from the given file. @@ -21,8 +34,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 calculate_checksum(path) diff --git a/src/simdb/cli/commands/simulation.py b/src/simdb/cli/commands/simulation.py index 59bd0bdd..54f4c4da 100644 --- a/src/simdb/cli/commands/simulation.py +++ b/src/simdb/cli/commands/simulation.py @@ -1,9 +1,10 @@ import contextlib import sys +import time import urllib.parse from itertools import chain from pathlib import Path -from typing import Any, List, Optional, Tuple, Type +from typing import Any, List, Optional, Tuple import appdirs import click @@ -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 @@ -200,19 +202,123 @@ def simulation_ingest(config: Config, manifest_file: str, alias: str): click.echo("ALIAS: " + simulation.alias + "\nUUID: " + str(simulation.uuid)) -def n_required_args_adaptor(n) -> Type[click.Command]: - class NRequiredArgs(click.Command): - NArgs = n +class OptionalRemoteCommand(click.Command): + """A command declared as `[REMOTE] ARG...` whose REMOTE may be left out.""" - def parse_args(self, ctx, args): - if len(args) == self.NArgs: - args.insert(0, "") - super().parse_args(ctx, args) + def parse_args(self, ctx, args): + arguments = [p for p in self.get_params(ctx) if isinstance(p, click.Argument)] + if self._count_values_given(ctx, args, arguments) < len(arguments): + args = ["", *args] + super().parse_args(ctx, args) - return NRequiredArgs + def _count_values_given(self, ctx, args, arguments) -> int: + """Count how many of the ARGUMENTS the command line provides a value for.""" + values = self.make_parser(ctx).parse_args(list(args))[0] + return sum(1 for argument in arguments if values.get(argument.name) is not None) -@simulation.command("push", cls=n_required_args_adaptor(1)) +def _prepare_simulation( + config: Config, api: RemoteAPI, sim_id: str, replaces: Optional[str] +) -> Simulation: + """Look up the local simulation SIM_ID and validate it against the remote + schemas.""" + db = get_local_db(config) + + simulation = db.get_simulation(sim_id) + if simulation is None: + raise click.ClickException(f"Failed to find simulation: {sim_id}") + + if replaces: + simulation.set_meta("replaces", replaces) + + schemas = api.get_validation_schemas() + try: + for schema in schemas: + Validator(schema).validate(simulation) + except ValidationError as err: + raise click.ClickException(f"Simulation does not validate: {err}") from err + + return simulation + + +def _wait_for_ingestion(api: RemoteAPI, sim_id: str, timeout: float) -> IngestionStatus: + """Poll the remote until the ingestion of SIM_ID reaches a terminal state. + + Reports every status change and returns the terminal status. + + :raise click.ClickException: if the remote cannot be reached, reports an + unknown status, or TIMEOUT seconds pass before + the ingestion finishes. + """ + max_consecutive_failures = 5 + poll_interval = 1.0 + + click.echo("Waiting for ingestion to complete...", nl=False) + last_status = None + consecutive_failures = 0 + deadline = time.monotonic() + timeout + while True: + try: + status = api.get_ingestion_status(sim_id) + except RemoteError as err: + # The remote rejected the request, so retrying will not help. + click.echo() + raise click.ClickException( + f"Failed to check ingestion status: {err}" + ) from err + except Exception as err: + # Tolerate transient errors: the ingestion continues server-side + consecutive_failures += 1 + if consecutive_failures >= max_consecutive_failures: + click.echo() + raise click.ClickException( + f"Failed to check ingestion status " + f"{consecutive_failures} times in a row: {err}" + ) from err + if time.monotonic() >= deadline: + click.echo() + raise click.ClickException( + f"Timed out after {timeout:g}s waiting for ingestion to " + f"complete (last status: {last_status})" + ) from err + time.sleep(poll_interval) + continue + consecutive_failures = 0 + + try: + ingestion_status = IngestionStatus(status) + except ValueError as err: + click.echo() + raise click.ClickException( + f"Remote reported an unknown ingestion status: {status}" + ) from err + + if status != last_status: + if last_status is not None: + click.echo(f" -> {status}", nl=False) + else: + click.echo(f" {status}", nl=False) + last_status = status + + if ingestion_status.is_terminal(): + click.echo() + return ingestion_status + + if time.monotonic() >= deadline: + click.echo() + raise click.ClickException( + f"Timed out after {timeout:g}s waiting for ingestion to complete " + f"(last status: {status})" + ) + + time.sleep(poll_interval) + + +@simulation.command( + "push_local", + cls=OptionalRemoteCommand, + short_help="Register a simulation whose files the REMOTE copies itself.", +) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") @@ -224,7 +330,14 @@ def parse_args(self, ctx, args): is_flag=True, help="Add the current user as a watcher of the simulation.", ) -def simulation_push( +@click.option( + "--timeout", + type=float, + default=600.0, + show_default=True, + help="Maximum number of seconds to wait for ingestion to complete.", +) +def simulation_push_local( config: Config, remote: Optional[str], sim_id: str, @@ -232,32 +345,70 @@ def simulation_push( password: Optional[str], replaces: Optional[str], add_watcher: bool, + timeout: float, ): - """Push the simulation with the given SIM_ID (UUID or alias) to the REMOTE.""" + """Register the simulation with the given SIM_ID (UUID or alias) on the REMOTE. + + Only the metadata is sent: the REMOTE copies the simulation files itself from + the recorded paths, which must therefore be reachable from the remote. + Waits for the remote ingestion to reach a terminal state, or until --timeout + seconds have passed. + """ api = RemoteAPI(remote, username, password, config) - db = get_local_db(config) + simulation = _prepare_simulation(config, api, sim_id, replaces) - simulation = db.get_simulation(sim_id) - if simulation is None: - raise click.ClickException(f"Failed to find simulation: {sim_id}") + api.push_local_simulation(simulation, add_watcher=add_watcher) - if replaces: - simulation.set_meta("replaces", replaces) + status = _wait_for_ingestion(api, simulation.uuid.hex, timeout) + if status is not IngestionStatus.COMPLETED: + raise click.ClickException( + f"Simulation ingestion failed with status: {status.value}" + ) - schemas = api.get_validation_schemas() - try: - for schema in schemas: - Validator(schema).validate(simulation) - except ValidationError as err: - raise click.ClickException(f"Simulation does not validate: {err}") from err + click.echo(f"Successfully pushed simulation {simulation.uuid}") + + +@simulation.command( + "push", + cls=OptionalRemoteCommand, + short_help="Upload a simulation and its files to the REMOTE.", +) +@pass_config +@click.argument("remote", required=False) +@click.argument("sim_id") +@click.option("--username", help="Username used to authenticate with the remote.") +@click.option("--password", help="Password used to authenticate with the remote.") +@click.option("--replaces", help="SIM_ID of simulation to deprecate and replace.") +@click.option( + "--add-watcher", + is_flag=True, + help="Add the current user as a watcher of the simulation.", +) +def simulation_push( + config: Config, + remote: Optional[str], + sim_id: str, + username: Optional[str], + password: Optional[str], + replaces: Optional[str], + add_watcher: bool, +): + """Push the simulation with the given SIM_ID (UUID or alias) to the REMOTE. + + Both the metadata and the simulation files are uploaded over HTTP. Use + push_local instead when the REMOTE can read the files itself. + """ + + api = RemoteAPI(remote, username, password, config) + simulation = _prepare_simulation(config, api, sim_id, replaces) api.push_simulation(simulation, out_stream=sys.stdout, add_watcher=add_watcher) click.echo(f"Successfully pushed simulation {simulation.uuid}") -@simulation.command("pull", cls=n_required_args_adaptor(2)) +@simulation.command("pull", cls=OptionalRemoteCommand) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") @@ -380,7 +531,7 @@ def simulation_query( ) -@simulation.command("data", cls=n_required_args_adaptor(2)) +@simulation.command("data", cls=OptionalRemoteCommand) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") @@ -448,7 +599,7 @@ def simulation_data( print_quantity(coord, label=f"coord {coord['name']}", show_stats=False) -@simulation.command("validate", cls=n_required_args_adaptor(1)) +@simulation.command("validate", cls=OptionalRemoteCommand) @pass_config @click.argument("remote", required=False) @click.argument("sim_id") diff --git a/src/simdb/cli/remote_api.py b/src/simdb/cli/remote_api.py index 97348cdb..4fee42eb 100644 --- a/src/simdb/cli/remote_api.py +++ b/src/simdb/cli/remote_api.py @@ -24,6 +24,7 @@ Optional, Tuple, Union, + cast, ) from urllib.parse import urlparse @@ -33,11 +34,13 @@ from requests.auth import AuthBase from semantic_version import Version +from simdb.checksum import calculate_checksum from simdb.config import Config from simdb.database.models import Simulation from simdb.imas.utils import SimDBUrl, 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 +238,87 @@ def _get_paths(file: "File") -> Iterable[Path]: return imas_files(file.uri) +def _find_partition_for_file( + file: Path, partitions: Dict[str, str] +) -> Tuple[str, Path]: + # Match the partition with the longest root so that a catch-all mapping + # (e.g. "/") does not shadow more specific partitions. + best: Optional[Tuple[str, Path]] = None + best_depth = -1 + for partition, path in partitions.items(): + root = Path(path) + try: + relative = file.relative_to(root) + except ValueError: + continue + depth = len(root.parts) + if depth > best_depth: + best = (partition, relative) + best_depth = depth + if best is None: + raise APIError( + f"File {file} is not located under any configured partition " + f"(configured partitions: {', '.join(partitions) or 'none'})" + ) + return best + + +def _file_data_for_partition( + file: FileData, source: Path, partitions: Dict[str, str], keep_uuid: bool = True +) -> FileData: + """Copy FILE with its URI rewritten relative to the partition holding SOURCE.""" + partition, partition_path = _find_partition_for_file(source, partitions) + new_uri = SimDBUrl.build(scheme=partition, path=partition_path.as_posix()) + update: Dict[str, Any] = { + "uri": new_uri.encoded_string(), + "checksum": calculate_checksum(source), + } + if not keep_uuid: + update["uuid"] = uuid.uuid1() + return file.model_copy(update=update) + + +def _source_files(file: FileData) -> List[Path]: + """Return the local files that FILE refers to, expanding directories.""" + file_uri = SimDBUrl(file.uri) + if file_uri.path is None: + raise APIError(f"File URI has no path: {file.uri}") + + if file_uri.scheme == "imas": + try: + sources = sorted(imas_files(file_uri)) + except ValueError as err: + raise APIError(f"Failed to list IMAS files of {file.uri}: {err}") from err + if not sources: + raise APIError(f"IMAS URI does not contain any files: {file.uri}") + return sources + + file_path = Path(file_uri.path) + if not file_path.is_dir(): + return [file_path] + + sources = [] + for sub_file in sorted(file_path.iterdir()): + if sub_file.is_dir(): + raise APIError(f"Nested directory found in {file_path}: {sub_file.name}") + sources.append(sub_file) + return sources + + +def _expand_directories( + files: Iterable[FileData], partitions: Dict[str, str] +) -> List[FileData]: + new_file_list = [] + for file in files: + sources = _source_files(file) + keep_uuid = len(sources) == 1 + for source in sources: + new_file_list.append( + _file_data_for_partition(file, source, partitions, keep_uuid=keep_uuid) + ) + return new_file_list + + class RemoteAPI: """ Class to represent connection to remote API. @@ -855,6 +939,32 @@ 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 = cast( + Dict[str, str], self._config.get_section("partition", default={}) + ) + 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")) + + @versioned_method("v1.3") + @try_request + def get_ingestion_status(self, sim_id: str) -> str: + res = self.get(f"simulation/status/{sim_id}") + return res.json()["status"] + @versioned_method("v1.2", "v1.3") @try_request def push_simulation( diff --git a/src/simdb/imas/utils.py b/src/simdb/imas/utils.py index 73a4c73e..700d886b 100644 --- a/src/simdb/imas/utils.py +++ b/src/simdb/imas/utils.py @@ -311,6 +311,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/workers/tasks.py b/src/simdb/workers/tasks.py index 97d9213d..9c7609e2 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 calculate_checksum 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") @@ -133,14 +109,6 @@ 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 _get_imas_identifier_path(path: Path) -> Path: if path.suffix == ".nc": return path @@ -153,7 +121,7 @@ def _create_file_from_data( uri = SimDBUrl(data.uri) path = _resolve_uri_to_path(uri, config) - checksum = _calculate_checksum(path) + checksum = calculate_checksum(path) if data.checksum != checksum: raise ValueError("Hash of file does not match provided checksum") @@ -182,7 +150,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) + checksum = calculate_checksum(path) if file_data.checksum != checksum: raise ValueError("Hash of file does not match provided checksum") file = File.from_data_model(file_data) diff --git a/tests/cli/test_cli_simulation_command.py b/tests/cli/test_cli_simulation_command.py index 0120fc02..1f81def4 100644 --- a/tests/cli/test_cli_simulation_command.py +++ b/tests/cli/test_cli_simulation_command.py @@ -1,9 +1,13 @@ from unittest import mock +from uuid import uuid1 +import pytest from click.testing import CliRunner from utils import config_test_file +from simdb.cli.remote_api import FailedConnection, RemoteError from simdb.cli.simdb import cli +from simdb.enums import IngestionStatus @mock.patch("simdb.database.get_local_db") @@ -85,3 +89,199 @@ def test_simulation_validate_command(remote_api, get_local_db): runner = CliRunner() result = runner.invoke(cli, [f"--config-file={config_file}", "simulation"]) assert result.exception is None + + +@pytest.mark.parametrize( + "subcommand, trailing_args", + ( + ("push", ()), + ("push_local", ()), + ("pull", ("directory",)), + ("data", ("ids_path",)), + ("validate", ()), + ), +) +@pytest.mark.parametrize( + "options", (("--username", "bob"), ()), ids=("username", "none") +) +@pytest.mark.parametrize( + "options_first", (True, False), ids=("options-first", "options-last") +) +@pytest.mark.parametrize("remote", (("test",), ()), ids=("test", "default")) +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_optional_remote_argument( + get_local_db, remote_api, remote, options_first, options, subcommand, trailing_args +): + """REMOTE may be left out, wherever the options appear on the command line.""" + config_file = config_test_file() + # push_local polls until the ingestion reaches a terminal state. + remote_api.return_value.get_ingestion_status.return_value = ( + IngestionStatus.COMPLETED.value + ) + arguments = (*remote, "sim_id", *trailing_args) + argv = (*options, *arguments) if options_first else (*arguments, *options) + + runner = CliRunner() + result = runner.invoke( + cli, [f"--config-file={config_file}", "simulation", subcommand, *argv] + ) + + assert remote_api.called, result.output + used_remote, used_username = remote_api.call_args.args[:2] + assert used_remote == (remote[0] if remote else "") + assert used_username == ("bob" if options else None) + + +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_command(get_local_db, remote_api): + """push_local waits for the remote ingestion to complete.""" + config_file = config_test_file() + remote_api.return_value.get_ingestion_status.return_value = ( + IngestionStatus.COMPLETED.value + ) + simulation = get_local_db.return_value.get_simulation.return_value + + runner = CliRunner() + result = runner.invoke( + cli, [f"--config-file={config_file}", "simulation", "push_local", "sim_id"] + ) + + assert result.exception is None, result.output + remote_api.return_value.push_local_simulation.assert_called_once_with( + simulation, add_watcher=False + ) + assert f"Successfully pushed simulation {simulation.uuid}" in result.output + + +@pytest.mark.parametrize( + "status", (IngestionStatus.COPY_FAILED, IngestionStatus.VALIDATION_FAILED) +) +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_failed_ingestion(get_local_db, remote_api, status): + """push_local fails when the remote ingestion does.""" + config_file = config_test_file() + remote_api.return_value.get_ingestion_status.return_value = status.value + + runner = CliRunner() + result = runner.invoke( + cli, [f"--config-file={config_file}", "simulation", "push_local", "sim_id"] + ) + + assert result.exit_code != 0 + assert f"Simulation ingestion failed with status: {status.value}" in result.output + + +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_timeout(get_local_db, remote_api): + """push_local stops waiting once the --timeout has passed.""" + config_file = config_test_file() + remote_api.return_value.get_ingestion_status.return_value = ( + IngestionStatus.COPYING.value + ) + argv = ("push_local", "sim_id", "--timeout=0") + + runner = CliRunner() + result = runner.invoke(cli, [f"--config-file={config_file}", "simulation", *argv]) + + assert result.exit_code != 0 + assert "Timed out after 0s waiting for ingestion to complete" in result.output + assert "last status: COPYING" in result.output + + +def _invoke_push_local(remote_api, *extra_args): + config_file = config_test_file() + runner = CliRunner() + return runner.invoke( + cli, + [ + f"--config-file={config_file}", + "simulation", + "push_local", + "sim_id", + *extra_args, + ], + ) + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_reports_status_changes(get_local_db, remote_api, sleep): + """push_local reports every ingestion status the remote goes through.""" + sim_uuid = uuid1() + get_local_db.return_value.get_simulation.return_value.uuid = sim_uuid + remote_api.return_value.get_ingestion_status.side_effect = [ + IngestionStatus.QUEUED.value, + IngestionStatus.COPYING.value, + IngestionStatus.COPYING.value, + IngestionStatus.COMPLETED.value, + ] + + result = _invoke_push_local(remote_api) + + assert result.exit_code == 0, result.output + assert "QUEUED -> COPYING -> COMPLETED" in result.output + assert f"Successfully pushed simulation {sim_uuid}" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_unknown_status(get_local_db, remote_api, sleep): + """push_local does not wait for a status it cannot recognise.""" + remote_api.return_value.get_ingestion_status.return_value = "TELEPORTING" + + result = _invoke_push_local(remote_api) + + assert result.exit_code != 0 + assert "Remote reported an unknown ingestion status: TELEPORTING" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_transient_errors(get_local_db, remote_api, sleep): + """The ingestion continues server-side, so a failed status check is not fatal.""" + remote_api.return_value.get_ingestion_status.side_effect = [ + FailedConnection("connection reset"), + FailedConnection("connection reset"), + IngestionStatus.COMPLETED.value, + ] + + result = _invoke_push_local(remote_api) + + assert result.exit_code == 0, result.output + assert "COMPLETED" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_persistent_errors(get_local_db, remote_api, sleep): + """push_local gives up once the status check keeps failing.""" + remote_api.return_value.get_ingestion_status.side_effect = FailedConnection("down") + + result = _invoke_push_local(remote_api) + + assert result.exit_code != 0 + assert "Failed to check ingestion status 5 times in a row: down" in result.output + + +@mock.patch("simdb.cli.commands.simulation.time.sleep") +@mock.patch("simdb.cli.commands.simulation.RemoteAPI") +@mock.patch("simdb.cli.commands.simulation.get_local_db") +def test_simulation_push_local_rejected_status_check(get_local_db, remote_api, sleep): + """A request the remote rejects will not start working when repeated.""" + remote_api.return_value.get_ingestion_status.side_effect = RemoteError( + "Simulation not found" + ) + + result = _invoke_push_local(remote_api) + + assert result.exit_code != 0 + assert "Failed to check ingestion status: Simulation not found" in result.output + assert remote_api.return_value.get_ingestion_status.call_count == 1 diff --git a/tests/cli/test_remote_api_push_local.py b/tests/cli/test_remote_api_push_local.py new file mode 100644 index 00000000..fcdb215f --- /dev/null +++ b/tests/cli/test_remote_api_push_local.py @@ -0,0 +1,108 @@ +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +from simdb.checksum import calculate_checksum +from simdb.cli.remote_api import ( + APIError, + _expand_directories, + _find_partition_for_file, +) +from simdb.remote.models import FileData + + +def _file_data(uri: str, file_type: str = "FILE") -> FileData: + return FileData( + type=file_type, + uri=uri, + checksum="stale", + datetime=datetime(2026, 1, 1, tzinfo=timezone.utc), + ) + + +def test_find_partition_prefers_the_deepest_root(): + """A catch-all partition does not shadow a more specific one.""" + partitions = {"root": "/", "data": "/mnt/data"} + + assert _find_partition_for_file(Path("/mnt/data/run/x.txt"), partitions) == ( + "data", + Path("run/x.txt"), + ) + assert _find_partition_for_file(Path("/sdcc/run/x.txt"), partitions) == ( + "root", + Path("sdcc/run/x.txt"), + ) + + +def test_find_partition_without_a_match(): + with pytest.raises(APIError, match="configured partitions: data"): + _find_partition_for_file(Path("/elsewhere/x.txt"), {"data": "/mnt/data"}) + + +def test_expand_directories_rewrites_the_uri_of_a_single_file(tmp_path: Path): + source = tmp_path / "run" / "x.txt" + source.parent.mkdir() + source.write_text("contents") + file = _file_data(f"file:{source}") + + expanded = _expand_directories([file], {"data": str(tmp_path)}) + + assert len(expanded) == 1 + assert expanded[0].uri == "data:run/x.txt" + assert expanded[0].checksum == calculate_checksum(source) + # A file that maps onto a single source keeps its identity. + assert expanded[0].uuid == file.uuid + + +def test_expand_directories_expands_a_directory(tmp_path: Path): + directory = tmp_path / "run" + directory.mkdir() + for name in ("b.txt", "a.txt"): + (directory / name).write_text(name) + file = _file_data(f"file:{directory}") + + expanded = _expand_directories([file], {"data": str(tmp_path)}) + + assert [f.uri for f in expanded] == ["data:run/a.txt", "data:run/b.txt"] + # Each file needs its own identity to be stored on the remote. + assert len({f.uuid for f in expanded}) == 2 + + +def test_expand_directories_rejects_nested_directories(tmp_path: Path): + directory = tmp_path / "run" + (directory / "nested").mkdir(parents=True) + file = _file_data(f"file:{directory}") + + with pytest.raises(APIError, match="Nested directory found"): + _expand_directories([file], {"data": str(tmp_path)}) + + +def test_expand_directories_only_lists_the_files_of_an_imas_backend(tmp_path: Path): + directory = tmp_path / "run" + directory.mkdir() + for name in ("master.h5", "equilibrium.h5", "notes.txt"): + (directory / name).write_text(name) + file = _file_data(f"imas:hdf5?path={directory}", file_type="IMAS") + + expanded = _expand_directories([file], {"data": str(tmp_path)}) + + assert [f.uri for f in expanded] == [ + "data:run/equilibrium.h5", + "data:run/master.h5", + ] + + +def test_expand_directories_reports_an_unusable_imas_uri(tmp_path: Path): + file = _file_data(f"imas:hdf5?path={tmp_path / 'missing'}", file_type="IMAS") + + with pytest.raises(APIError, match="Failed to list IMAS files"): + _expand_directories([file], {"data": str(tmp_path)}) + + +def test_expand_directories_without_configured_partitions(tmp_path: Path): + source = tmp_path / "x.txt" + source.write_text("contents") + + with pytest.raises(APIError, match="configured partitions: none"): + _expand_directories([_file_data(f"file:{source}")], {}) diff --git a/tests/remote/api/v1.3/test_simulations3.py b/tests/remote/api/v1.3/test_simulations3.py index 9fb6d9b6..dc8edb06 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 calculate_checksum 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 = calculate_checksum(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..26b6cabf 100644 --- a/tests/workers/test_tasks.py +++ b/tests/workers/test_tasks.py @@ -4,13 +4,13 @@ import pytest +from simdb.checksum import calculate_checksum 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, _copy_files, _create_file_from_data, _get_imas_identifier_path, @@ -189,7 +189,7 @@ def test_copy_files_task_copies_inputs_and_marks_copied(task_environment): input_files = [ _make_file_data( - f"data:/{source_file.name}", checksum=_calculate_checksum(source_file) + f"data:/{source_file.name}", checksum=calculate_checksum(source_file) ) ] diff --git a/uv.lock b/uv.lock index 4130318a..581d0cf7 100644 --- a/uv.lock +++ b/uv.lock @@ -1266,7 +1266,7 @@ requires-dist = [ { name = "imas-simdb", extras = ["imas-validator", "postgres", "server"], marker = "extra == 'all'" }, { name = "imas-validator", marker = "extra == 'imas-validator'", specifier = ">=1.0.0" }, { name = "myst-parser", marker = "extra == 'build-docs'", specifier = ">=0.18.0" }, - { name = "netcdf4", specifier = ">=1.5" }, + { name = "netcdf4", specifier = ">=1.7.2" }, { name = "numpy", specifier = ">=1.14" }, { name = "plotext", specifier = "==5.3.2" }, { name = "psycopg2-binary", marker = "extra == 'postgres'", specifier = ">=2.8.0" },