From 487cfa4fbc36ad0573b0317c728ad134b1f0f419 Mon Sep 17 00:00:00 2001 From: Simon Pinches Date: Thu, 18 Jun 2026 15:55:44 +0200 Subject: [PATCH 1/3] Sort IMAS file globs to make checksums platform-independent imas.checksum.checksum() feeds the files returned by imas_files() into a single running SHA1 hash in iteration order. imas_files() built the HDF5 and ASCII file lists from Path.glob(), which returns entries in an arbitrary, filesystem-dependent order. As a result the same byte-identical IMAS data could hash to different checksums on Windows vs Linux. Sort the glob results explicitly by file name so the iteration order is deterministic and identical across platforms. Sort by p.name rather than relying on Path comparison, which is itself platform-dependent (Windows folds case, Linux does not). Add regression tests for imas_files ordering. --- src/simdb/imas/utils.py | 4 +-- tests/test_imas_utils.py | 53 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 tests/test_imas_utils.py diff --git a/src/simdb/imas/utils.py b/src/simdb/imas/utils.py index 7b9ff234..6e02c92b 100644 --- a/src/simdb/imas/utils.py +++ b/src/simdb/imas/utils.py @@ -332,7 +332,7 @@ def imas_files(uri: SimDBUrl) -> List[Path]: path = _get_path(uri) if backend == "hdf5": - return [p.absolute() for p in path.glob("*.h5")] + return [p.absolute() for p in sorted(path.glob("*.h5"), key=lambda p: p.name)] elif backend == "mdsplus": return [ path / "ids_001.characteristics", @@ -340,7 +340,7 @@ def imas_files(uri: SimDBUrl) -> List[Path]: path / "ids_001.tree", ] elif backend == "ascii": - return [p.absolute() for p in path.glob("*.ids")] + return [p.absolute() for p in sorted(path.glob("*.ids"), key=lambda p: p.name)] else: raise ValueError(f"Unknown IMAS backend {backend}") diff --git a/tests/test_imas_utils.py b/tests/test_imas_utils.py new file mode 100644 index 00000000..bd048ebf --- /dev/null +++ b/tests/test_imas_utils.py @@ -0,0 +1,53 @@ +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + +from simdb.imas.utils import imas_files +from simdb.uri import URI + + +class ImasFilesTests(unittest.TestCase): + """Tests for simdb.imas.utils.imas_files. + + The checksum is computed by feeding the files into a single running hash + in the order imas_files returns them, so that order must be deterministic + and identical across platforms. Path.glob() does not sort, so imas_files + sorts explicitly by file name. See utils.imas_files / imas.checksum.checksum. + """ + + def _make_files(self, directory, names): + # Create files in an order that does not match the expected sorted order + for name in names: + (Path(directory) / name).write_bytes(b"") + + def test_hdf5_files_sorted_by_name(self): + names = [ + "equilibrium.h5", + "core_profiles.h5", + "master.h5", + "summary.h5", + ] + with TemporaryDirectory() as tmp: + self._make_files(tmp, names) + uri = URI(f"imas:hdf5?path={tmp}") + result = [p.name for p in imas_files(uri)] + self.assertEqual(result, sorted(names)) + + def test_ascii_files_sorted_by_name(self): + names = ["equilibrium.ids", "core_profiles.ids", "summary.ids"] + with TemporaryDirectory() as tmp: + self._make_files(tmp, names) + uri = URI(f"imas:ascii?path={tmp}") + result = [p.name for p in imas_files(uri)] + self.assertEqual(result, sorted(names)) + + def test_hdf5_files_returns_absolute_paths(self): + with TemporaryDirectory() as tmp: + self._make_files(tmp, ["core_profiles.h5"]) + uri = URI(f"imas:hdf5?path={tmp}") + result = imas_files(uri) + self.assertTrue(all(p.is_absolute() for p in result)) + + +if __name__ == "__main__": + unittest.main() From 508f114ee2887dad96f492cf51946e575bb43468 Mon Sep 17 00:00:00 2001 From: Yannick de Jong Date: Fri, 17 Jul 2026 10:59:01 +0200 Subject: [PATCH 2/3] Online migrations --- .../c3a1f0b9d4e2_widen_checksum_column.py | 48 ++++++ ...4b2e6f1a7c3_add_online_migrations_table.py | 37 +++++ src/simdb/checksum.py | 44 +++++- src/simdb/cli/commands/simulation.py | 3 +- src/simdb/cli/remote_api.py | 3 +- src/simdb/database/database.py | 18 +++ src/simdb/database/models/__init__.py | 10 +- src/simdb/database/models/file.py | 2 +- src/simdb/database/models/online_migration.py | 26 ++++ src/simdb/imas/checksum.py | 3 +- src/simdb/remote/apis/files.py | 44 +++++- src/simdb/workers/migrations.py | 141 ++++++++++++++++++ src/simdb/workers/tasks.py | 35 ++++- tests/remote/api/test_files.py | 23 ++- tests/remote/api/test_staging.py | 3 +- tests/test_imas_utils.py | 3 +- tests/workers/test_migrations.py | 91 +++++++++++ tests/workers/test_tasks.py | 16 +- 18 files changed, 533 insertions(+), 17 deletions(-) create mode 100644 alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py create mode 100644 alembic/versions/d4b2e6f1a7c3_add_online_migrations_table.py create mode 100644 src/simdb/database/models/online_migration.py create mode 100644 src/simdb/workers/migrations.py create mode 100644 tests/workers/test_migrations.py diff --git a/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py b/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py new file mode 100644 index 00000000..3c20ec74 --- /dev/null +++ b/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py @@ -0,0 +1,48 @@ +"""Widen files.checksum column for algorithm-prefixed checksums + +Checksums are now stored as ``:`` (e.g. ``sha1:...``). +Widen the column from 64 to 128 characters so the prefix fits today and leaves +room for longer digests (e.g. ``sha256:``) in the future. + +The actual re-hashing of existing values is done by the online (data) migration +``recalculate_checksums`` in :mod:`simdb.workers.migrations`, not here. + +Revision ID: c3a1f0b9d4e2 +Revises: b2c52ee8ff12 +Create Date: 2026-07-17 00:00:00.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "c3a1f0b9d4e2" +down_revision: Union[str, Sequence[str], None] = "b2c52ee8ff12" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + with op.batch_alter_table("files", schema=None) as batch_op: + batch_op.alter_column( + "checksum", + existing_type=sa.String(length=64), + type_=sa.String(length=128), + existing_nullable=True, + ) + + +def downgrade() -> None: + """Downgrade schema.""" + with op.batch_alter_table("files", schema=None) as batch_op: + batch_op.alter_column( + "checksum", + existing_type=sa.String(length=128), + type_=sa.String(length=64), + existing_nullable=True, + ) diff --git a/alembic/versions/d4b2e6f1a7c3_add_online_migrations_table.py b/alembic/versions/d4b2e6f1a7c3_add_online_migrations_table.py new file mode 100644 index 00000000..f60bc3d9 --- /dev/null +++ b/alembic/versions/d4b2e6f1a7c3_add_online_migrations_table.py @@ -0,0 +1,37 @@ +"""Add online_migrations tracking table + +Records which online (data) migrations have been applied, so the runner in +:mod:`simdb.workers.migrations` can skip migrations that have already run. + +Revision ID: d4b2e6f1a7c3 +Revises: c3a1f0b9d4e2 +Create Date: 2026-07-17 00:00:01.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "d4b2e6f1a7c3" +down_revision: Union[str, Sequence[str], None] = "c3a1f0b9d4e2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + op.create_table( + "online_migrations", + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("applied_at", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("name"), + ) + + +def downgrade() -> None: + """Downgrade schema.""" + op.drop_table("online_migrations") diff --git a/src/simdb/checksum.py b/src/simdb/checksum.py index 99aef230..fc3eaf14 100644 --- a/src/simdb/checksum.py +++ b/src/simdb/checksum.py @@ -3,12 +3,52 @@ from simdb.imas.utils import SimDBUrl +#: Algorithm used to generate checksums. Prepended to every stored checksum as a +#: ``:`` prefix so the encoding is self-describing. +CHECKSUM_ALGORITHM = "sha1" + + +def format_checksum(hexdigest: str, algorithm: str = CHECKSUM_ALGORITHM) -> str: + """Prefix a raw hex digest with its algorithm, e.g. ``sha1:2fd4e1c6...``. + + :param hexdigest: the hex representation of the digest + :param algorithm: the algorithm that produced the digest + :return: the algorithm-prefixed checksum string + """ + return f"{algorithm}:{hexdigest}" + + +def is_prefixed(checksum: str) -> bool: + """Return whether a checksum already carries an ``:`` prefix.""" + return bool(checksum) and ":" in checksum + + +def strip_checksum(checksum: str) -> str: + """Return the bare hex digest, dropping any ``:`` prefix. + + Legacy (pre-prefix) checksums are returned unchanged. Used to serialize + checksums on the wire for API versions that predate the prefix. + """ + if not checksum: + return checksum + return checksum.split(":", 1)[1] if ":" in checksum else checksum + + +def checksums_match(a: str, b: str) -> bool: + """Compare two checksums ignoring any algorithm prefix on either side. + + This keeps validation working across the prefix change: a legacy bare-hex + checksum (e.g. from an older client) is considered equal to its prefixed + form (``sha1:``). + """ + return strip_checksum(a) == strip_checksum(b) + def sha1_checksum(uri: SimDBUrl) -> str: """Generate a SHA1 checksum from the given file. :param uri: the URI of the file to checksum - :return: a string containing the hex representation of the computed SHA1 checksum + :return: the algorithm-prefixed checksum (``sha1:``) """ if uri.scheme != "file": raise ValueError(f"invalid scheme for file checksum: {uri.scheme}") @@ -25,4 +65,4 @@ def sha1_checksum(uri: SimDBUrl) -> str: with path.open("rb") as file: for chunk in iter(lambda: file.read(4096), b""): sha1.update(chunk) - return sha1.hexdigest() + return format_checksum(sha1.hexdigest()) diff --git a/src/simdb/cli/commands/simulation.py b/src/simdb/cli/commands/simulation.py index 59bd0bdd..dcf0024f 100644 --- a/src/simdb/cli/commands/simulation.py +++ b/src/simdb/cli/commands/simulation.py @@ -9,6 +9,7 @@ import click from rich.prompt import Confirm +from simdb.checksum import checksums_match from simdb.cli.manifest import Manifest from simdb.cli.remote_api import RemoteAPI, RemoteError from simdb.config.config import Config @@ -479,7 +480,7 @@ def simulation_validate( # Pass config and ids_list parameters current_checksum = file.generate_checksum(config, ids_list) - if current_checksum != file.checksum: + if not checksums_match(current_checksum, file.checksum): raise ValidationError( f"Checksum mismatch for file {file.uri}. " f"Expected: {file.checksum}, Got: {current_checksum}" diff --git a/src/simdb/cli/remote_api.py b/src/simdb/cli/remote_api.py index 99f87363..8cb46b0d 100644 --- a/src/simdb/cli/remote_api.py +++ b/src/simdb/cli/remote_api.py @@ -33,6 +33,7 @@ from requests.auth import AuthBase from semantic_version import Version +from simdb.checksum import checksums_match from simdb.config import Config from simdb.database.models import Simulation from simdb.imas.utils import SimDBUrl, imas_files @@ -1084,7 +1085,7 @@ def _pull_file( ) print("\r", file=out_stream, end="", flush=True) - if sha1.hexdigest() != checksum: + if not checksums_match(sha1.hexdigest(), checksum): raise APIError(f"Checksum failed for file {from_path}") @versioned_method("v1.2", "v1.3") diff --git a/src/simdb/database/database.py b/src/simdb/database/database.py index f8523a11..eb974099 100644 --- a/src/simdb/database/database.py +++ b/src/simdb/database/database.py @@ -933,9 +933,27 @@ def get_local_db(config: Config) -> Database: run_migrations(database.engine) else: raise e + # With the schema at head, apply any pending online (data) migrations. This + # runs on every open, but is a cheap no-op once the data is up to date. + _run_online_migrations(database, config) return database +def _run_online_migrations(database: Database, config: Config) -> None: + """Apply online (data) migrations to the local database. + + Recalculates any checksums still stored in the legacy (bare-hex) format. + Imported lazily to avoid a circular import: ``simdb.workers`` imports + ``tasks``, which imports ``get_db`` from this module. + """ + from simdb.workers.migrations import run_online_migrations # noqa: PLC0415 + + results = run_online_migrations(database, config) + changed = sum(results.values()) + if changed: + print(f"Recalculated {changed} local checksum(s).") + + def get_db(config: Config) -> Database: db_type = config.get_option("database.type") if db_type == "postgres": diff --git a/src/simdb/database/models/__init__.py b/src/simdb/database/models/__init__.py index ea062b17..9f2b5cd3 100644 --- a/src/simdb/database/models/__init__.py +++ b/src/simdb/database/models/__init__.py @@ -1,7 +1,15 @@ from .base import Base from .file import File from .metadata import MetaData +from .online_migration import OnlineMigrationHistory from .simulation import Simulation from .watcher import Watcher -__all__ = ["Base", "File", "MetaData", "Simulation", "Watcher"] +__all__ = [ + "Base", + "File", + "MetaData", + "OnlineMigrationHistory", + "Simulation", + "Watcher", +] diff --git a/src/simdb/database/models/file.py b/src/simdb/database/models/file.py index 202b94dc..20ba4b29 100644 --- a/src/simdb/database/models/file.py +++ b/src/simdb/database/models/file.py @@ -30,7 +30,7 @@ class File(Base): id = Column(sql_types.Integer, primary_key=True) uuid = Column(UUID, nullable=False, unique=True, index=True) uri: SimDBUrl = Column(URI(1024), nullable=True) - checksum = Column(sql_types.String(64), nullable=True) + checksum = Column(sql_types.String(128), nullable=True) type = Column(sql_types.Enum(DataType), nullable=True) datetime = Column(sql_types.DateTime, nullable=False) diff --git a/src/simdb/database/models/online_migration.py b/src/simdb/database/models/online_migration.py new file mode 100644 index 00000000..ef71930d --- /dev/null +++ b/src/simdb/database/models/online_migration.py @@ -0,0 +1,26 @@ +from datetime import datetime + +from sqlalchemy import Column +from sqlalchemy import types as sql_types + +from .base import Base + + +class OnlineMigrationHistory(Base): + """Record of an applied online (data) migration. + + Online migrations transform live data (see :mod:`simdb.workers.migrations`). + One row is written per migration once it has completed successfully, which + lets the runner skip migrations that have already run. + """ + + __tablename__ = "online_migrations" + name = Column(sql_types.String(255), primary_key=True) + applied_at = Column(sql_types.DateTime, nullable=False, default=datetime.now) + + def __init__(self, name: str, applied_at: datetime) -> None: + self.name = name + self.applied_at = applied_at + + def __str__(self) -> str: + return f"{self.name} (applied {self.applied_at})" diff --git a/src/simdb/imas/checksum.py b/src/simdb/imas/checksum.py index d9d403ef..7769b9af 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 format_checksum from simdb.imas.utils import SimDBUrl from .utils import imas_files, list_idss, open_imas @@ -27,4 +28,4 @@ def checksum(uri: SimDBUrl, ids_list: list) -> str: continue for chunk in iter(lambda: file.read(4096), b""): sha1.update(chunk) - return sha1.hexdigest() + return format_checksum(sha1.hexdigest()) diff --git a/src/simdb/remote/apis/files.py b/src/simdb/remote/apis/files.py index cd7cad5e..c3293121 100644 --- a/src/simdb/remote/apis/files.py +++ b/src/simdb/remote/apis/files.py @@ -1,4 +1,5 @@ import gzip +import re import uuid from pathlib import Path from typing import Dict, Iterable, List, Optional @@ -8,7 +9,7 @@ from flask_restx import Namespace, Resource from werkzeug.datastructures import FileStorage -from simdb.checksum import sha1_checksum +from simdb.checksum import checksums_match, sha1_checksum, strip_checksum from simdb.cli.manifest import DataType from simdb.database import DatabaseError, models from simdb.imas.checksum import checksum as imas_checksum @@ -28,6 +29,30 @@ api = Namespace("files", path="/") +#: First API version whose wire format carries the ``:`` checksum +#: prefix. Older clients expect a bare hex digest and compare it exactly, so we +#: strip the prefix from responses served under earlier versions. +_PREFIX_WIRE_MIN_VERSION = (1, 3) + + +def _request_api_version() -> tuple: + """Return the API version of the current request as a ``(major, minor)`` tuple. + + Derived from the Flask blueprint name (e.g. ``api_v1_2`` -> ``(1, 2)``). + Defaults to the newest behaviour when it cannot be determined. + """ + match = re.match(r"api_v(\d+)(?:_(\d+))?", request.blueprint or "") + if not match: + return _PREFIX_WIRE_MIN_VERSION + return (int(match.group(1)), int(match.group(2) or 0)) + + +def _wire_checksum(checksum: str) -> str: + """Serialize a stored checksum for the wire, stripping the prefix pre-v1.3.""" + if checksum and _request_api_version() < _PREFIX_WIRE_MIN_VERSION: + return strip_checksum(checksum) + return checksum + def _verify_file( sim_uuid: uuid.UUID, @@ -50,7 +75,7 @@ def _verify_file( if not path.exists(): raise ValueError(f"file {path} does not exist") checksum = sha1_checksum(SimDBUrl.build(scheme="file", path=path.as_posix())) - if sim_file.checksum != checksum: + if not checksums_match(sim_file.checksum, checksum): raise ValueError(f"checksum failed for file {sim_file!r}") elif sim_file.type == DataType.IMAS: uri = sim_file.uri @@ -69,7 +94,7 @@ def _verify_file( scheme=uri.scheme, path=uri.path, query=f"path={path_value}" ) checksum = imas_checksum(new_uri, ids_list or []) - if sim_file.checksum != checksum: + if not checksums_match(sim_file.checksum, checksum): raise ValueError(f"checksum failed for simulation {sim_file.uri}") @@ -178,7 +203,12 @@ class FileList(Resource): @pydantic_validate(api) def get(self, user: User) -> FileDataList: files = current_app.db.list_files() - return FileDataList.model_validate([file.to_model() for file in files]) + models_ = [] + for file in files: + model = file.to_model() + model.checksum = _wire_checksum(model.checksum) + models_.append(model) + return FileDataList.model_validate(models_) @requires_auth() def post(self, user: User): @@ -198,7 +228,11 @@ class File(Resource): @pydantic_validate(api) def get(self, file_uuid: str, user: Optional[User] = None) -> FileGetDataResponse: file = current_app.db.get_file(file_uuid) - return file.to_model_with_path() + response = file.to_model_with_path() + response.checksum = _wire_checksum(response.checksum) + for file_info in response.files: + file_info.checksum = _wire_checksum(file_info.checksum) + return response @api.route("/file/download/") diff --git a/src/simdb/workers/migrations.py b/src/simdb/workers/migrations.py new file mode 100644 index 00000000..c2de1b71 --- /dev/null +++ b/src/simdb/workers/migrations.py @@ -0,0 +1,141 @@ +"""Online (data) migrations. + +Unlike Alembic migrations, which change the database *schema* offline, these +migrations transform live *data* while the system is online. They are run +automatically: + +* on the server, by :func:`simdb.workers.tasks.run_online_migrations_task` + which is queued when a Celery worker starts; +* for the local SQLite database, by ``get_local_db`` after the schema is + brought up to date. + +Which migrations have run is tracked in the ``online_migrations`` table (see +:class:`simdb.database.models.OnlineMigrationHistory`): a migration is recorded +once it completes successfully and is skipped on subsequent runs. A migration +that only partially completes (e.g. a file was temporarily unavailable) is not +recorded, so it is retried on the next run -- migrations must therefore be +idempotent. + +To add a migration, write a function ``def my_migration(database, config) -> +MigrationResult`` and append an :class:`OnlineMigration` to +:data:`ONLINE_MIGRATIONS` with a unique, stable ``name``. +""" + +import logging +from dataclasses import dataclass +from datetime import datetime +from typing import Callable, Dict, Set + +from simdb.config import Config +from simdb.database.database import Database +from simdb.database.models import File, OnlineMigrationHistory + +logger = logging.getLogger(__name__) + + +@dataclass +class MigrationResult: + """Outcome of a single online migration run. + + :param changed: number of rows changed during this run. + :param complete: whether the migration finished with nothing left to do. + When ``False`` the migration is not recorded and will run again. + """ + + changed: int + complete: bool = True + + +@dataclass(frozen=True) +class OnlineMigration: + """A single idempotent data migration.""" + + name: str + run: Callable[[Database, Config], MigrationResult] + + +def _recalculate_checksums(database: Database, config: Config) -> MigrationResult: + """Recalculate and prefix checksums that predate the ``sha1:`` format. + + Any file whose stored checksum lacks an ``:`` prefix is + re-hashed from disk (picking up the platform-independent, sorted-glob + computation) and stored in the new ``sha1:`` form. Files that cannot be + hashed right now are left unchanged and reported as incomplete so the + migration is retried later. + """ + # Only files whose checksum lacks an ":" prefix still need + # migrating. Empty checksums (checksum generation disabled) are left alone. + unmigrated = ( + database.session.query(File) + .filter(File.checksum.isnot(None)) + .filter(File.checksum != "") + .filter(File.checksum.notlike("%:%")) + .all() + ) + + updated = 0 + failed = 0 + for file in unmigrated: + try: + file.checksum = file.generate_checksum(config, []) + except Exception: + failed += 1 + logger.exception( + "Could not recalculate checksum for file %s; leaving it unchanged", + file.uri, + ) + continue + updated += 1 + if updated: + database.session.commit() + if failed: + logger.warning( + "Checksum recalculation left %d file(s) unmigrated (see errors above); " + "they will be retried on the next run", + failed, + ) + return MigrationResult(changed=updated, complete=(failed == 0)) + + +ONLINE_MIGRATIONS = [ + OnlineMigration(name="recalculate_checksums", run=_recalculate_checksums), +] + + +def _applied_migration_names(database: Database) -> Set[str]: + return {row[0] for row in database.session.query(OnlineMigrationHistory.name).all()} + + +def _record_migration(database: Database, name: str) -> None: + database.session.add(OnlineMigrationHistory(name=name, applied_at=datetime.now())) + database.session.commit() + + +def run_online_migrations(database: Database, config: Config) -> Dict[str, int]: + """Run every not-yet-applied online migration in order. + + :return: a mapping of migration name to rows changed, for migrations that + actually ran this time (already-applied migrations are omitted). + """ + applied = _applied_migration_names(database) + results: Dict[str, int] = {} + for migration in ONLINE_MIGRATIONS: + if migration.name in applied: + continue + result = migration.run(database, config) + results[migration.name] = result.changed + if result.complete: + _record_migration(database, migration.name) + logger.info( + "Online migration %r applied (%d row(s) changed)", + migration.name, + result.changed, + ) + else: + logger.warning( + "Online migration %r did not fully complete (%d row(s) changed); " + "it will run again next time", + migration.name, + result.changed, + ) + return results diff --git a/src/simdb/workers/tasks.py b/src/simdb/workers/tasks.py index 97d9213d..57ba85d6 100644 --- a/src/simdb/workers/tasks.py +++ b/src/simdb/workers/tasks.py @@ -8,8 +8,10 @@ from typing import Iterable, List from uuid import UUID +from celery.signals import worker_ready from pydantic import AnyUrl +from simdb.checksum import format_checksum, is_prefixed from simdb.config import Config from simdb.database.database import get_db from simdb.database.models import File @@ -18,6 +20,7 @@ from simdb.imas.utils import SimDBUrl from simdb.remote.models import FileData, FileDataList from simdb.workers.celery import celery_app +from simdb.workers.migrations import run_online_migrations logger = logging.getLogger(__name__) @@ -138,7 +141,7 @@ def _calculate_checksum(path: Path) -> str: with path.open("rb") as f: for chunk in iter(lambda: f.read(4096), b""): sha1.update(chunk) - return sha1.hexdigest() + return format_checksum(sha1.hexdigest()) def _get_imas_identifier_path(path: Path) -> Path: @@ -154,6 +157,8 @@ def _create_file_from_data( path = _resolve_uri_to_path(uri, config) checksum = _calculate_checksum(path) + if not is_prefixed(data.checksum): + raise ValueError("Checksum must include an algorithm prefix (e.g. 'sha1:...')") if data.checksum != checksum: raise ValueError("Hash of file does not match provided checksum") @@ -183,6 +188,10 @@ def _create_files_from_data_list( file = _create_file_from_data(file_data, config, imas_path) else: checksum = _calculate_checksum(path) + if not is_prefixed(file_data.checksum): + raise ValueError( + "Checksum must include an algorithm prefix (e.g. 'sha1:...')" + ) if file_data.checksum != checksum: raise ValueError("Hash of file does not match provided checksum") file = File.from_data_model(file_data) @@ -314,3 +323,27 @@ def fail_stale_ingestions_task() -> dict: return {"failed": failed} finally: database.close() + + +@celery_app.task +def run_online_migrations_task() -> dict: + """Run all pending online (data) migrations against the server database. + + Automatically queued when a Celery worker becomes ready. Every migration is + idempotent, so running this repeatedly (e.g. once per worker) is safe. + """ + config = Config() + config.load() + database = get_db(config) + + try: + results = run_online_migrations(database, config) + return {"status": "completed", "migrations": results} + finally: + database.close() + + +@worker_ready.connect +def _queue_online_migrations(sender=None, **_kwargs) -> None: + """Queue the online migrations once a worker is ready to process them.""" + run_online_migrations_task.delay() diff --git a/tests/remote/api/test_files.py b/tests/remote/api/test_files.py index e4e6e328..414327bb 100644 --- a/tests/remote/api/test_files.py +++ b/tests/remote/api/test_files.py @@ -3,8 +3,10 @@ import io import json import tarfile +import types from datetime import datetime, timezone from pathlib import Path +from unittest import mock import pytest from conftest import ( @@ -13,6 +15,8 @@ post_simulation, ) +import simdb.remote.apis.files as files_api +from simdb.checksum import format_checksum from simdb.cli.manifest import DataType from simdb.json import CustomEncoder from simdb.remote.models import ( @@ -36,7 +40,7 @@ def create_simulation_with_file( for i in range(0, len(file_content), chunk_size) ] num_chunks = len(chunks) - test_checksum = hashlib.sha1(file_content).hexdigest() + test_checksum = format_checksum(hashlib.sha1(file_content).hexdigest()) simulation_data = generate_simulation_data( alias=alias, @@ -189,3 +193,20 @@ def test_download_file(client): assert rv.status_code == 200 assert rv.data == file_content + + +@pytest.mark.parametrize( + "blueprint,expected", + [ + ("api_v1", "deadbeef"), + ("api_v1_1", "deadbeef"), + ("api_v1_2", "deadbeef"), + ("api_v1_3", "sha1:deadbeef"), + ("", "sha1:deadbeef"), # unknown -> newest behaviour keeps the prefix + ], +) +def test_wire_checksum_strips_prefix_before_v1_3(blueprint, expected): + with mock.patch.object( + files_api, "request", types.SimpleNamespace(blueprint=blueprint) + ): + assert files_api._wire_checksum("sha1:deadbeef") == expected diff --git a/tests/remote/api/test_staging.py b/tests/remote/api/test_staging.py index 9e795c23..213dd4ad 100644 --- a/tests/remote/api/test_staging.py +++ b/tests/remote/api/test_staging.py @@ -4,6 +4,7 @@ from conftest import HEADERS, generate_simulation_data, post_simulation +from simdb.checksum import format_checksum from simdb.remote.models import FileData, StagingDirectoryResponse @@ -27,7 +28,7 @@ def test_get_staging_dir_for_simulation_with_uuid(client): def test_create_simulation_from_staging_dir(client_copy_files): file_data = b"test_data" - checksum = hashlib.sha1(file_data).hexdigest() + checksum = format_checksum(hashlib.sha1(file_data).hexdigest()) simulation_data = generate_simulation_data( alias="test-simulation", inputs=[ diff --git a/tests/test_imas_utils.py b/tests/test_imas_utils.py index bd048ebf..7d69244e 100644 --- a/tests/test_imas_utils.py +++ b/tests/test_imas_utils.py @@ -2,9 +2,10 @@ from pathlib import Path from tempfile import TemporaryDirectory -from simdb.imas.utils import imas_files from simdb.uri import URI +from simdb.imas.utils import imas_files + class ImasFilesTests(unittest.TestCase): """Tests for simdb.imas.utils.imas_files. diff --git a/tests/workers/test_migrations.py b/tests/workers/test_migrations.py new file mode 100644 index 00000000..7f2baf97 --- /dev/null +++ b/tests/workers/test_migrations.py @@ -0,0 +1,91 @@ +import hashlib +import tempfile +import uuid +from datetime import datetime, timezone + +import pytest + +from simdb.checksum import is_prefixed, sha1_checksum +from simdb.cli.manifest import DataType +from simdb.database import Database +from simdb.database.models import Base, File, OnlineMigrationHistory +from simdb.imas.utils import SimDBUrl +from simdb.workers.migrations import run_online_migrations + + +@pytest.fixture +def db(): + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + db_file = f.name + database = Database(Database.DBMS.SQLITE, file=db_file) + Base.metadata.create_all(database.engine) + yield database + database.close() + + +def _add_file(db, uri: SimDBUrl, checksum: str) -> File: + file = File(DataType.FILE, uri, perform_integrity_check=False) + file.uuid = uuid.uuid1() + file.checksum = checksum + file.datetime = datetime.now(timezone.utc) + db.session.add(file) + db.session.commit() + return file + + +def _applied(db): + return {m.name for m in db.session.query(OnlineMigrationHistory).all()} + + +def test_recalculates_and_prefixes_legacy_checksums(db, tmp_path): + target = tmp_path / "data.txt" + target.write_text("hello world") + uri = SimDBUrl.build(scheme="file", path=target.as_posix()) + + expected = sha1_checksum(uri) # sha1: + legacy = _add_file(db, uri, hashlib.sha1(b"hello world").hexdigest()) + already = _add_file(db, uri, expected) + disabled = _add_file(db, uri, "") + + results = run_online_migrations(db, config=None) + + assert results["recalculate_checksums"] == 1 + db.session.refresh(legacy) + assert legacy.checksum == expected + assert is_prefixed(legacy.checksum) + assert already.checksum == expected # untouched + assert disabled.checksum == "" # untouched + assert "recalculate_checksums" in _applied(db) + + +def test_completed_migration_is_recorded_and_skipped(db, tmp_path): + target = tmp_path / "data.txt" + target.write_text("content") + uri = SimDBUrl.build(scheme="file", path=target.as_posix()) + _add_file(db, uri, hashlib.sha1(b"content").hexdigest()) + + first = run_online_migrations(db, config=None) + assert first["recalculate_checksums"] == 1 + assert "recalculate_checksums" in _applied(db) + + # Already recorded -> not run again (omitted from the results mapping). + second = run_online_migrations(db, config=None) + assert "recalculate_checksums" not in second + + +def test_incomplete_migration_is_not_recorded_and_retries(db, tmp_path): + # File points at a path that does not exist, so it cannot be hashed. + missing_uri = SimDBUrl.build(scheme="file", path=str(tmp_path / "gone.txt")) + file = _add_file(db, missing_uri, "deadbeef") + + results = run_online_migrations(db, config=None) + assert results["recalculate_checksums"] == 0 + assert file.checksum == "deadbeef" # left for retry + assert "recalculate_checksums" not in _applied(db) + + # Once the file becomes available, a later run migrates it and records it. + (tmp_path / "gone.txt").write_text("now here") + results = run_online_migrations(db, config=None) + assert results["recalculate_checksums"] == 1 + assert is_prefixed(file.checksum) + assert "recalculate_checksums" in _applied(db) diff --git a/tests/workers/test_tasks.py b/tests/workers/test_tasks.py index 6d411be4..4ae3f961 100644 --- a/tests/workers/test_tasks.py +++ b/tests/workers/test_tasks.py @@ -143,12 +143,26 @@ def test_create_file_from_data_raises_on_checksum_mismatch(tmp_path): data_file = partition_path / "testfile.txt" data_file.write_text("content") - file_data = _make_file_data("data:testfile.txt", checksum="wrong_checksum") + file_data = _make_file_data("data:testfile.txt", checksum="sha1:wrong_checksum") with pytest.raises(ValueError, match="Hash of file does not match"): _create_file_from_data(file_data, config, data_file) +def test_create_file_from_data_raises_on_unprefixed_checksum(tmp_path): + config = Config() + partition_path = tmp_path / "partition_data" + partition_path.mkdir() + config.set_option("partition.data", str(partition_path)) + data_file = partition_path / "testfile.txt" + data_file.write_text("content") + + file_data = _make_file_data("data:testfile.txt", checksum="deadbeef") + + with pytest.raises(ValueError, match="algorithm prefix"): + _create_file_from_data(file_data, config, data_file) + + @pytest.fixture def task_environment(tmp_path): """Set up Config, mocked DB, and directory layout for copy_files_task tests.""" From 4b2b841530d52973c2c8dfc34eb0940971e31180 Mon Sep 17 00:00:00 2001 From: Yannick de Jong Date: Thu, 30 Jul 2026 11:42:43 +0200 Subject: [PATCH 3/3] Fix tests --- .../c3a1f0b9d4e2_widen_checksum_column.py | 4 +- tests/test_imas_utils.py | 96 +++++++++---------- 2 files changed, 45 insertions(+), 55 deletions(-) diff --git a/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py b/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py index 3c20ec74..d6fb633f 100644 --- a/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py +++ b/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py @@ -8,7 +8,7 @@ ``recalculate_checksums`` in :mod:`simdb.workers.migrations`, not here. Revision ID: c3a1f0b9d4e2 -Revises: b2c52ee8ff12 +Revises: 6fb9b8fbac38 Create Date: 2026-07-17 00:00:00.000000 """ @@ -21,7 +21,7 @@ # revision identifiers, used by Alembic. revision: str = "c3a1f0b9d4e2" -down_revision: Union[str, Sequence[str], None] = "b2c52ee8ff12" +down_revision: Union[str, Sequence[str], None] = "6fb9b8fbac38" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None diff --git a/tests/test_imas_utils.py b/tests/test_imas_utils.py index 7d69244e..fee52faa 100644 --- a/tests/test_imas_utils.py +++ b/tests/test_imas_utils.py @@ -1,54 +1,44 @@ -import unittest from pathlib import Path -from tempfile import TemporaryDirectory - -from simdb.uri import URI - -from simdb.imas.utils import imas_files - - -class ImasFilesTests(unittest.TestCase): - """Tests for simdb.imas.utils.imas_files. - - The checksum is computed by feeding the files into a single running hash - in the order imas_files returns them, so that order must be deterministic - and identical across platforms. Path.glob() does not sort, so imas_files - sorts explicitly by file name. See utils.imas_files / imas.checksum.checksum. - """ - - def _make_files(self, directory, names): - # Create files in an order that does not match the expected sorted order - for name in names: - (Path(directory) / name).write_bytes(b"") - - def test_hdf5_files_sorted_by_name(self): - names = [ - "equilibrium.h5", - "core_profiles.h5", - "master.h5", - "summary.h5", - ] - with TemporaryDirectory() as tmp: - self._make_files(tmp, names) - uri = URI(f"imas:hdf5?path={tmp}") - result = [p.name for p in imas_files(uri)] - self.assertEqual(result, sorted(names)) - - def test_ascii_files_sorted_by_name(self): - names = ["equilibrium.ids", "core_profiles.ids", "summary.ids"] - with TemporaryDirectory() as tmp: - self._make_files(tmp, names) - uri = URI(f"imas:ascii?path={tmp}") - result = [p.name for p in imas_files(uri)] - self.assertEqual(result, sorted(names)) - - def test_hdf5_files_returns_absolute_paths(self): - with TemporaryDirectory() as tmp: - self._make_files(tmp, ["core_profiles.h5"]) - uri = URI(f"imas:hdf5?path={tmp}") - result = imas_files(uri) - self.assertTrue(all(p.is_absolute() for p in result)) - - -if __name__ == "__main__": - unittest.main() + +from simdb.imas.utils import SimDBUrl, imas_files + +# Tests for simdb.imas.utils.imas_files. +# +# The checksum is computed by feeding the files into a single running hash +# in the order imas_files returns them, so that order must be deterministic +# and identical across platforms. Path.glob() does not sort, so imas_files +# sorts explicitly by file name. See utils.imas_files / imas.checksum.checksum. + + +def _make_files(directory, names): + # Create files in an order that does not match the expected sorted order + for name in names: + (Path(directory) / name).write_bytes(b"") + + +def test_hdf5_files_sorted_by_name(tmp_path): + names = [ + "equilibrium.h5", + "core_profiles.h5", + "master.h5", + "summary.h5", + ] + _make_files(tmp_path, names) + uri = SimDBUrl(f"imas:hdf5?path={tmp_path}") + result = [p.name for p in imas_files(uri)] + assert result == sorted(names) + + +def test_ascii_files_sorted_by_name(tmp_path): + names = ["equilibrium.ids", "core_profiles.ids", "summary.ids"] + _make_files(tmp_path, names) + uri = SimDBUrl(f"imas:ascii?path={tmp_path}") + result = [p.name for p in imas_files(uri)] + assert result == sorted(names) + + +def test_hdf5_files_returns_absolute_paths(tmp_path): + _make_files(tmp_path, ["core_profiles.h5"]) + uri = SimDBUrl(f"imas:hdf5?path={tmp_path}") + result = imas_files(uri) + assert all(p.is_absolute() for p in result)