diff --git a/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py b/alembic/versions/c3a1f0b9d4e2_widen_checksum_column.py new file mode 100644 index 00000000..d6fb633f --- /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: 6fb9b8fbac38 +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] = "6fb9b8fbac38" +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/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/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 new file mode 100644 index 00000000..fee52faa --- /dev/null +++ b/tests/test_imas_utils.py @@ -0,0 +1,44 @@ +from pathlib import Path + +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) 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."""