diff --git a/docs/changes/newsfragments/8219.new b/docs/changes/newsfragments/8219.new new file mode 100644 index 000000000000..c9564aa5db18 --- /dev/null +++ b/docs/changes/newsfragments/8219.new @@ -0,0 +1,20 @@ +Added **split raw data storage** for the DataSet: an opt-in results +backend selected via the ``dataset.raw_data_backend`` config option +(set it to ``"sqlite_per_dataset_db"``) that writes raw measurement +data into individual per-dataset SQLite files while keeping all metadata +in the main database. This prevents the main DB file from growing +excessively large and allows users to have a single main SQLite QCoDeS +database on their PC as opposed to many that grow infinitely. When +selected, no results table is created in the main database at all - it +holds only metadata - and all public ``DataSet`` APIs continue to work +unchanged. Backend-specific settings (such as the per-dataset file +folder ``raw_data_path``) live under ``dataset.raw_data_backend_config``. + +Management helpers included: + +- ``update_raw_data_paths()``: update stored paths after raw data files + have been moved to a new folder. +- ``purge_orphaned_datasets()``: find and remove dataset records whose + raw data files no longer exist on disk (``dry_run=True`` by default). +- ``cleanup_datasets()``: remove datasets (DB records and raw data files) + by age, sample name, or file size (``dry_run=True`` by default). diff --git a/docs/dataset/dataset_design.rst b/docs/dataset/dataset_design.rst index f44795586083..a9e00082d3a1 100644 --- a/docs/dataset/dataset_design.rst +++ b/docs/dataset/dataset_design.rst @@ -75,3 +75,51 @@ We note that the dataset currently exclusively supports storing data in an SQLite database. This is not an intrinsic limitation of the dataset and measurement layer. It is possible that at a future state support for writing to a different backend will be added. + +.. _sec:design_split_storage: + +Split Raw Data Storage +====================== + +As the main SQLite database grows with many datasets, managing the database +file can become inconvenient due to the file size. To address this, +QCoDeS supports an optional **split raw data storage** mode (see +:ref:`sec:intro_split_raw_data` for user-facing details). + +From a design perspective, this feature adds a pluggable *results backend* +inside the ``DataSet`` class without changing any public interfaces: + +- A ``ResultsBackend`` strategy (in ``qcodes.dataset._results_backend``) + encapsulates where and how a dataset's results table is stored. The default + ``MainDatabaseResultsBackend`` keeps results in the main database; + ``SeparateSqliteFileResultsBackend`` writes them to a per-dataset SQLite file. + Backends are registered by their ``backend_name`` (the ``dataset.raw_data_backend`` + config value); the backend is selected in ``DataSet.__init__`` from that config + for new runs, and from the run's recorded state for existing ones. Adding a + backend is a matter of registering a new ``ResultsBackend`` subclass and a + ``raw_data_backend`` / ``raw_data_backend_config`` entry. +- A ``_results_conn`` property returns the backend's connection -- the main + database connection by default, or a per-dataset raw data connection when + results are stored separately. +- Write paths (``add_results``, ``_BackgroundWriter``) and read paths + (``get_parameter_data``, ``DataSetCacheWithDBBackend``, ``number_of_results``, + ``__len__``) all go through the backend / ``_results_conn``. +- The per-dataset SQLite file is a lightweight database containing only the + results table and numpy type adapters -- no QCoDeS metadata schema. +- When raw data storage is enabled, **no results table is created in the main + database** at all -- only the run metadata is stored there. This mirrors how + ``DataSetInMem`` records runs. A run is identified as a split-storage dataset + by the ``raw_data_db_path`` column in the ``runs`` table (recorded at dataset + creation), which is used both to reconnect to the raw data file and to + distinguish such runs from ``DataSetInMem`` runs (which also have no results + table). This column is an internal storage detail and is kept out of the + user-facing metadata. +- Subscriber triggers (used for real-time data callbacks) are created on the + results connection. Because the results table only exists once the dataset is + started, subscriptions requested before then are deferred and materialised at + start time. + +The implementation is contained in ``qcodes.dataset._raw_data_storage`` (helper +functions) and a handful of additions to ``qcodes.dataset.data_set`` (routing +logic). The ``Measurement`` context manager, ``DataSaver``, and all export +functions work without modification. diff --git a/docs/dataset/introduction.rst b/docs/dataset/introduction.rst index b5de7a35e4a3..f4474192a4cd 100644 --- a/docs/dataset/introduction.rst +++ b/docs/dataset/introduction.rst @@ -75,3 +75,52 @@ For dataset operations, QCoDeS provides functions for: - **Exporting datasets**: :doc:`Exporting data to other file formats <../examples/DataSet/Exporting-data-to-other-file-formats>` - **Extracting runs between databases**: :doc:`Extracting runs from one DB file to another <../examples/DataSet/Extracting-runs-from-one-DB-file-to-another>` and :func:`qcodes.dataset.extract_runs_into_db` - **Bulk export and metadata-only databases**: :func:`qcodes.dataset.export_datasets_and_create_metadata_db` for creating lightweight metadata-only databases while exporting all data to NetCDF files + +.. _sec:intro_split_raw_data: + +Split Raw Data Storage +====================== + +By default, all measurement data (the results table rows) is stored in the same SQLite database alongside metadata such as experiments, runs, parameters and their dependencies. Over time, the main database file can grow very large, which can slow down operations like browsing experiments and make managing a large database file inconvenient. + +QCoDeS supports an optional **split raw data storage** mode in which the actual measurement data for each ``DataSet`` is written to an individual, per-dataset SQLite file while all metadata remains in the main database. Each per-dataset file is named after the dataset's GUID (e.g. ``.db``) and is stored in a configurable folder. + +Which backend stores a dataset's results is selected by the ``dataset.raw_data_backend`` configuration option, with backend-specific settings under ``dataset.raw_data_backend_config`` (see the :doc:`Configuring QCoDeS notebook <../examples/basic_examples/Configuring_QCoDeS>` for how to view and change QCoDeS configuration): + +- ``dataset.raw_data_backend`` (string, default ``"sqlite_main_db"``): the results backend. ``"sqlite_main_db"`` keeps results in the main database; ``"sqlite_per_dataset_db"`` enables split storage. +- ``dataset.raw_data_backend_config.sqlite_per_dataset_db.raw_data_path`` (string, default ``"{db_location}"``): the folder where per-dataset files are created. The ``{db_location}`` placeholder is expanded to a folder derived from the main database path (e.g. ``~/experiments.db`` becomes ``~/experiments_db/``). + +When ``"sqlite_per_dataset_db"`` is selected: + +- No results table is created in the main database; it stores only the run metadata, keeping it lightweight. +- All ``INSERT`` and ``SELECT`` operations on results data are transparently routed to the per-dataset file. +- The path to the per-dataset file is recorded in the ``runs`` table, so ``load_by_id`` and related loading functions automatically reconnect to the correct file. +- All public ``DataSet`` APIs (``get_parameter_data``, ``to_pandas_dataframe``, ``to_xarray_dataset``, ``cache``, ``export``, etc.) work identically whichever backend is selected. + +Example runtime configuration:: + + import qcodes as qc + + qc.config.dataset.raw_data_backend = "sqlite_per_dataset_db" + qc.config.dataset.raw_data_backend_config.sqlite_per_dataset_db.raw_data_path = "/data/raw_measurements/" + +If the per-dataset raw data files are moved to a different folder (e.g. during data migration or archival), the stored paths in the main database will become stale. Use the :func:`~qcodes.dataset.update_raw_data_paths` helper to update them:: + + from qcodes.dataset import update_raw_data_paths + + update_raw_data_paths( + db_path="/path/to/main_database.db", + new_raw_data_folder="/new/location/of/raw_files/" + ) + +This scans all datasets with the ``raw_data_db_path`` internal ``runs`` column set, checks whether the corresponding ``.db`` file exists in the new folder, and updates the stored path accordingly. + +Managing Datasets +----------------- + +QCoDeS provides helpers for managing per-dataset raw data files over time: + +- :func:`~qcodes.dataset.purge_orphaned_datasets` — remove dataset records whose raw data files no longer exist on disk. +- :func:`~qcodes.dataset.cleanup_datasets` — remove datasets (DB records and raw data files) by age, sample name, or file size. + +For usage examples and more details on database management, see the :doc:`Database notebook <../examples/DataSet/Database>`. diff --git a/docs/examples/DataSet/Database.ipynb b/docs/examples/DataSet/Database.ipynb index 9aeaa6bd818d..14f3267c23bd 100644 --- a/docs/examples/DataSet/Database.ipynb +++ b/docs/examples/DataSet/Database.ipynb @@ -167,6 +167,143 @@ "\n", "Moreover, we have also written an [example notebook](Extracting-runs-from-one-DB-file-to-another.ipynb) of transferring `DataSets` between database flies that may serve as a template for more complex data organization protocols." ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Split Raw Data Storage\n", + "\n", + "As the main database grows with many datasets, browsing experiments and loading metadata can become slower. QCoDeS supports an optional **split raw data storage** mode that writes the raw measurement data for each dataset into its own individual SQLite file, while keeping all metadata (experiments, runs, parameters, dependencies) in the main database.\n", + "\n", + "This keeps the main database lightweight and makes it faster to work with, while still allowing all existing `DataSet` APIs to function identically.\n", + "\n", + "### Configuration\n", + "\n", + "Split raw data storage is selected via the `dataset.raw_data_backend` option, with backend-specific settings under `dataset.raw_data_backend_config`:\n", + "\n", + "- `dataset.raw_data_backend` (string, default `\"sqlite_main_db\"`): the results backend. Set to `\"sqlite_per_dataset_db\"` to enable split storage.\n", + "- `dataset.raw_data_backend_config.sqlite_per_dataset_db.raw_data_path` (string, default `\"{db_location}\"`): the folder where per-dataset SQLite files are created. The `{db_location}` placeholder expands to a folder derived from the main database path (e.g. `~/experiments.db` becomes `~/experiments_db/`).\n", + "\n", + "You can enable it at runtime:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Enable split raw data storage\n", + "qc.config.dataset.raw_data_backend = \"sqlite_per_dataset_db\"\n", + "\n", + "# Optionally set a custom path for per-dataset files\n", + "qc.config.dataset.raw_data_backend_config.sqlite_per_dataset_db.raw_data_path = (\n", + " \"/data/raw_measurements/\"\n", + ")\n", + "\n", + "# Or use the default which derives from the main DB location:\n", + "# qc.config.dataset.raw_data_backend_config.sqlite_per_dataset_db.raw_data_path = \"{db_location}\"\n", + "# e.g. ~/experiments.db -> ~/experiments_db/" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To make split storage the default for future sessions, set the same options through the `qc.config` Python interface shown above and persist your configuration - see [Configuring QCoDeS](../basic_examples/Configuring_QCoDeS.ipynb).\n", + "\n", + "### How It Works\n", + "\n", + "When split storage is enabled:\n", + "\n", + "1. When a measurement starts (`mark_started()`), a per-dataset SQLite file named `.db` is created in the configured folder.\n", + "2. All measurement data (results table rows) is written to this per-dataset file instead of the main database.\n", + "3. No results table is created in the main database for such datasets; only the run metadata (experiments, runs, parameters, dependencies) is kept there, keeping the main database small.\n", + "4. The path to the per-dataset file is saved in a dedicated internal `raw_data_db_path` column of the `runs` table (not exposed via `DataSet.metadata`), so `load_by_id()` and related functions automatically find and reconnect to the correct file.\n", + "5. All `DataSet` methods (`get_parameter_data`, `to_pandas_dataframe`, `to_xarray_dataset`, `cache`, `export`, etc.) work transparently with split storage.\n", + "\n", + "> **Note:** Datasets created with split storage enabled can always be loaded later, even if the configuration is changed back to the default, as long as the per-dataset files remain at their original paths.\n", + "\n", + "### Updating Paths After Moving Raw Data Files\n", + "\n", + "If you move the per-dataset raw data files to a different folder, the paths stored in the main database become stale. Use `update_raw_data_paths` to fix them:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# from qcodes.dataset import update_raw_data_paths\n", + "\n", + "# After moving raw data files to a new folder:\n", + "# update_raw_data_paths(\n", + "# db_path=\"/path/to/main_database.db\",\n", + "# new_raw_data_folder=\"/new/location/of/raw_files/\"\n", + "# )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The function scans all datasets that have the internal `raw_data_db_path` column set, checks whether the corresponding `.db` file exists in the new folder, and updates the stored path. Datasets whose files are not found in the new folder are skipped with a warning." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Purging Orphaned Datasets\n", + "\n", + "After archiving or deleting selected raw data files, some dataset records in the main database may reference files that no longer exist. Use `purge_orphaned_datasets` to clean up those orphaned records:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# from qcodes.dataset import purge_orphaned_datasets\n", + "\n", + "# Dry run (default) — see what would be removed:\n", + "# result = purge_orphaned_datasets(\"/path/to/main_database.db\")\n", + "# print(f\"{len(result.orphaned_datasets)} orphaned datasets found\")\n", + "\n", + "# Actually remove orphaned records:\n", + "# result = purge_orphaned_datasets(\"/path/to/main_database.db\", dry_run=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Cleaning Up Datasets by Criteria\n", + "\n", + "To free disk space by removing datasets (both raw data files and DB records) based on age, sample name, or file size, use `cleanup_datasets`. Criteria are combined with AND logic — a dataset must match **all** specified criteria to be selected." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# from qcodes.dataset import cleanup_datasets\n", + "\n", + "# Remove datasets older than 30 days (dry run first):\n", + "# result = cleanup_datasets(\"/path/to/db.db\", older_than_days=30)\n", + "# print(f\"{len(result.matching_datasets)} datasets would be removed\")\n", + "\n", + "# Remove datasets for a specific sample:\n", + "# result = cleanup_datasets(\"/path/to/db.db\", sample_name=\"old-sample\", dry_run=False)\n", + "\n", + "# Remove datasets with raw data files larger than 500 MB:\n", + "# result = cleanup_datasets(\"/path/to/db.db\", larger_than_mb=500, dry_run=False)" + ] } ], "metadata": { diff --git a/src/qcodes/configuration/qcodesrc.json b/src/qcodes/configuration/qcodesrc.json index 583f36bf8326..7f9083b1da1e 100644 --- a/src/qcodes/configuration/qcodesrc.json +++ b/src/qcodes/configuration/qcodesrc.json @@ -79,7 +79,13 @@ "export_chunked_export_of_large_files_enabled": false, "export_chunked_threshold": 1000, "in_memory_cache": true, - "load_from_exported_file": false + "load_from_exported_file": false, + "raw_data_backend": "sqlite_main_db", + "raw_data_backend_config": { + "sqlite_per_dataset_db": { + "raw_data_path": "{db_location}" + } + } }, "telemetry": { diff --git a/src/qcodes/configuration/qcodesrc_schema.json b/src/qcodes/configuration/qcodesrc_schema.json index ddd4929be22c..0ec9d64eaac8 100644 --- a/src/qcodes/configuration/qcodesrc_schema.json +++ b/src/qcodes/configuration/qcodesrc_schema.json @@ -382,6 +382,35 @@ "type": "boolean", "default": true, "description": "Should the data be cached in memory as it is measured. Useful to disable for large datasets to save on memory consumption." + }, + "raw_data_backend": { + "type": "string", + "enum": ["sqlite_main_db", "sqlite_per_dataset_db"], + "default": "sqlite_main_db", + "description": "Which backend stores dataset raw measurement data (the results tables). 'sqlite_main_db' (default) keeps the results in the main QCoDeS database. 'sqlite_per_dataset_db' writes each dataset's results to an individual per-dataset SQLite file while all metadata stays in the main database. Backend-specific settings live under raw_data_backend_config." + }, + "raw_data_backend_config": { + "type": "object", + "default": { + "sqlite_per_dataset_db": { + "raw_data_path": "{db_location}" + } + }, + "description": "Backend-specific configuration, keyed by the raw_data_backend name. Each backend reads its own entry; add a new key here (and a new raw_data_backend enum value) when adding a new backend.", + "properties": { + "sqlite_per_dataset_db": { + "type": "object", + "description": "Configuration for the 'sqlite_per_dataset_db' results backend.", + "properties": { + "raw_data_path": { + "type": "string", + "default": "{db_location}", + "description": "Path to the folder where per-dataset raw data SQLite files are stored. {db_location} is a directory in the same folder as the .db file with a matching name, e.g. for ~/experiments.db raw data files will be stored in ~/experiments_db/" + } + }, + "required": ["raw_data_path"] + } + } } }, "description": "Settings related to the DataSet and Measurement Context manager", diff --git a/src/qcodes/dataset/__init__.py b/src/qcodes/dataset/__init__.py index 0494069ab0b5..a44e101c6e9f 100644 --- a/src/qcodes/dataset/__init__.py +++ b/src/qcodes/dataset/__init__.py @@ -3,6 +3,11 @@ and from disk """ +from ._raw_data_storage import ( + cleanup_datasets, + purge_orphaned_datasets, + update_raw_data_paths, +) from .data_set import ( get_guids_by_run_spec, load_by_counter, @@ -86,6 +91,7 @@ "ThreadPoolParamsCaller", "TogetherSweep", "call_params_threaded", + "cleanup_datasets", "connect", "datasaver_builder", "do0d", @@ -121,6 +127,8 @@ "new_experiment", "plot_by_id", "plot_dataset", + "purge_orphaned_datasets", "reset_default_experiment_id", "rundescriber_from_json", + "update_raw_data_paths", ] diff --git a/src/qcodes/dataset/_raw_data_storage.py b/src/qcodes/dataset/_raw_data_storage.py new file mode 100644 index 000000000000..61a5168eaaa0 --- /dev/null +++ b/src/qcodes/dataset/_raw_data_storage.py @@ -0,0 +1,578 @@ +""" +Module for managing per-dataset raw data SQLite files. + +When ``dataset.raw_data_backend`` is set to ``"sqlite_per_dataset_db"``, +measurement data (results tables) are written to individual SQLite files +- one per dataset - instead of the main QCoDeS database file. All metadata +(runs, experiments, parameters) remains in the main database. + +The per-dataset files are stored in the folder given by +``dataset.raw_data_backend_config.sqlite_per_dataset_db.raw_data_path`` and are +named ``.db``. +""" + +from __future__ import annotations + +import logging +from contextlib import closing +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import TYPE_CHECKING, NamedTuple + +from tqdm.auto import tqdm + +import qcodes +from qcodes.dataset.export_config import _expand_export_path +from qcodes.dataset.sqlite.connection import AtomicConnection, atomic +from qcodes.dataset.sqlite.database import ( + _connect_to_sqlite_file, + connect, +) +from qcodes.dataset.sqlite.queries import ( + _create_run_table, + _remove_dataset_from_db, + get_datasets_with_raw_data_path, +) +from qcodes.dataset.sqlite.query_helpers import insert_column, is_column_in_table + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + from typing import Any + + from qcodes.parameters import ParamSpecBase + +log = logging.getLogger(__name__) + +_DATASET_CONFIG_SECTION = "dataset" +_RESULTS_BACKEND_KEY = "raw_data_backend" +_BACKEND_CONFIG_KEY = "raw_data_backend_config" +_RAW_DATA_PATH_KEY = "raw_data_path" + +#: ``raw_data_backend`` value keeping results in the main QCoDeS database. +MAIN_DB_BACKEND = "sqlite_main_db" +#: ``raw_data_backend`` value selecting the per-dataset SQLite-file backend. +PER_DATASET_DB_BACKEND = "sqlite_per_dataset_db" + + +def get_configured_results_backend_name() -> str: + """Return the results backend name from ``dataset.raw_data_backend``.""" + return qcodes.config[_DATASET_CONFIG_SECTION].get( + _RESULTS_BACKEND_KEY, MAIN_DB_BACKEND + ) + + +def get_results_backend_config(backend_name: str) -> Mapping[str, Any]: + """Return the backend-specific config for *backend_name*. + + Reads ``dataset.raw_data_backend_config.``; returns an empty + mapping when the backend has no configured settings. + """ + all_config = qcodes.config[_DATASET_CONFIG_SECTION].get(_BACKEND_CONFIG_KEY, {}) + return all_config.get(backend_name, {}) or {} + + +def is_raw_data_storage_enabled() -> bool: + """Return True if the per-dataset SQLite results backend is selected.""" + return get_configured_results_backend_name() == PER_DATASET_DB_BACKEND + + +def get_raw_data_folder(db_path: str | None = None) -> Path: + """Return the resolved folder path for raw data SQLite files. + + The path template comes from the ``sqlite_per_dataset_db`` backend config + and is expanded the same way as the export path. ``{db_location}`` is + resolved relative to *db_path* (the owning dataset's database) when given, + otherwise relative to the global ``core.db_location`` config. + """ + config = get_results_backend_config(PER_DATASET_DB_BACKEND) + raw_path_template: str = config.get(_RAW_DATA_PATH_KEY, "{db_location}") + return Path(_expand_export_path(raw_path_template, db_path)).expanduser().absolute() + + +def get_raw_data_db_path( + guid: str, folder: Path | None = None, db_path: str | None = None +) -> Path: + """Return the full path for a dataset's raw data SQLite file. + + Args: + guid: The GUID of the dataset. + folder: Override folder. If *None*, uses :func:`get_raw_data_folder`. + db_path: The owning dataset's database file, used to resolve + ``{db_location}`` when *folder* is not given. + + """ + if folder is None: + folder = get_raw_data_folder(db_path) + return folder / f"{guid}.db" + + +def connect_to_raw_data_db( + path: str | Path, + *, + read_only: bool = False, +) -> AtomicConnection: + """Open (or create) a lightweight SQLite connection for raw data. + + Unlike the main QCoDeS :func:`~qcodes.dataset.sqlite.database.connect`, + this does **not** create the full metadata schema (experiments, runs, ...). + It reuses the shared + :func:`~qcodes.dataset.sqlite.database._connect_to_sqlite_file` helper so + that the numpy/sqlite type adapters and connection settings are identical + to the main database connection. + + Args: + path: Path to the raw-data SQLite file. + read_only: Open the database in read-only mode. + + Returns: + An :class:`AtomicConnection` to the raw-data database. + + """ + return _connect_to_sqlite_file(path, read_only=read_only) + + +def create_raw_data_db( + path: str | Path, + table_name: str, + paramspecs: Sequence[ParamSpecBase], +) -> AtomicConnection: + """Create a per-dataset raw-data SQLite file with a results table. + + The file is created if it does not exist. The parent directory is + created if needed. + + Args: + path: Full path for the new SQLite file. + table_name: Name of the results table to create (matches the name + in the main database). + paramspecs: Parameter specifications describing the columns. + + Returns: + An :class:`AtomicConnection` to the newly created database. + + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + conn = connect_to_raw_data_db(path) + + # Create the empty results table, then add the parameter columns with + # insert_column - exactly as the main database does - so column names are + # quoted. ``_create_run_table`` with paramspecs would emit them verbatim via + # ``ParamSpecBase.sql_repr()``, which breaks for names that are SQL keywords + # (e.g. a parameter called ``from``). + _create_run_table(conn, table_name) + with atomic(conn) as aconn: + for spec in paramspecs: + insert_column(aconn, table_name, spec.name, spec.type) + + log.info( + "Created raw data database at %s with table %s", + path, + table_name, + ) + return conn + + +def update_raw_data_paths( + db_path: str | Path, + new_raw_data_folder: str | Path, +) -> list[tuple[int, str, str]]: + """Update raw data file paths in the main database after files have moved. + + Use this when per-dataset raw data files have been relocated to a new + folder but the main database still references the old paths. + + The function scans all runs that have a ``raw_data_db_path`` internal + ``runs`` column set, verifies that a file with the expected GUID-based name + exists in *new_raw_data_folder*, and updates the stored path in the + database. + + Args: + db_path: Path to the main QCoDeS database file. + new_raw_data_folder: The new folder where the per-dataset SQLite + files now reside. + + Returns: + A list of ``(run_id, old_path, new_path)`` tuples for every run + whose path was updated. + + Raises: + FileNotFoundError: If the main database file does not exist. + FileNotFoundError: If *new_raw_data_folder* does not exist. + + """ + db_path = Path(db_path) + # Resolve to an absolute path so the stored path does not depend on the + # process working directory when datasets are loaded later (newly created + # split datasets also store absolute paths). + new_raw_data_folder = Path(new_raw_data_folder).resolve() + + if not db_path.is_file(): + raise FileNotFoundError(f"Database file not found: {db_path}") + if not new_raw_data_folder.is_dir(): + raise FileNotFoundError(f"New raw data folder not found: {new_raw_data_folder}") + + updated: list[tuple[int, str, str]] = [] + + with closing(connect(str(db_path))) as conn: + if not is_column_in_table(conn, "runs", "raw_data_db_path"): + log.info( + "No raw_data_db_path column found in %s; nothing to update.", db_path + ) + return [] + + cursor = conn.execute( + "SELECT run_id, raw_data_db_path FROM runs " + "WHERE raw_data_db_path IS NOT NULL" + ) + rows = cursor.fetchall() + + for run_id, old_path_str in tqdm(rows, desc="Updating raw data paths"): + old_path = Path(old_path_str) + # The per-dataset file name is always .db — preserved on move + new_path = new_raw_data_folder / old_path.name + + if not new_path.is_file(): + log.warning( + "Run %d: expected raw data file %s not found in new folder;" + " skipping.", + run_id, + new_path, + ) + continue + + if str(new_path) == old_path_str: + continue # already correct + + new_path_str = str(new_path) + with atomic(conn) as aconn: + aconn.execute( + "UPDATE runs SET raw_data_db_path = ? WHERE run_id = ?", + (new_path_str, run_id), + ) + updated.append((run_id, old_path_str, new_path_str)) + log.debug( + "Run %d: updated raw_data_db_path from %s to %s", + run_id, + old_path_str, + new_path_str, + ) + + log.info("Updated %d raw data paths in %s", len(updated), db_path) + return updated + + +# --------------------------------------------------------------------------- +# Dataset management helpers +# --------------------------------------------------------------------------- + + +class DatasetInfo(NamedTuple): + """Summary information about a dataset in the main database.""" + + run_id: int + guid: str + experiment_name: str + sample_name: str + run_timestamp: float | None + completed_timestamp: float | None + result_table_name: str + raw_data_db_path: str | None + raw_data_size_bytes: int | None + + +@dataclass +class PurgeResult: + """Result of a purge_orphaned_datasets operation.""" + + total_datasets_with_raw_data: int + orphaned_datasets: list[DatasetInfo] + removed_datasets: list[DatasetInfo] + dry_run: bool + errors: list[tuple[int, Exception]] = field(default_factory=list) + + +@dataclass +class CleanupResult: + """Result of a cleanup_datasets operation.""" + + total_datasets_scanned: int + matching_datasets: list[DatasetInfo] + removed_datasets: list[DatasetInfo] + total_size_freed_bytes: int + dry_run: bool + errors: list[tuple[int, Exception]] = field(default_factory=list) + + +def _build_dataset_info_list( + conn: AtomicConnection, +) -> list[DatasetInfo]: + """Query datasets with raw data paths and enrich with file size info.""" + rows = get_datasets_with_raw_data_path(conn) + + datasets: list[DatasetInfo] = [] + for row in tqdm(rows, desc="Scanning datasets"): + raw_path = row.raw_data_db_path + raw_size: int | None = None + if raw_path and Path(raw_path).is_file(): + raw_size = Path(raw_path).stat().st_size + + datasets.append( + DatasetInfo( + run_id=row.run_id, + guid=row.guid, + experiment_name=row.experiment_name, + sample_name=row.sample_name, + run_timestamp=row.run_timestamp, + completed_timestamp=row.completed_timestamp, + result_table_name=row.result_table_name, + raw_data_db_path=raw_path, + raw_data_size_bytes=raw_size, + ) + ) + return datasets + + +def purge_orphaned_datasets( + db_path: str | Path, + *, + dry_run: bool = True, +) -> PurgeResult: + """Find and optionally remove dataset records whose raw data files are missing. + + When using split raw data storage, users may archive and delete + individual per-dataset SQLite files. This function identifies + datasets in the main database that reference raw data files which + no longer exist on disk, and optionally removes those dataset + records from the main database. + + Args: + db_path: Path to the main QCoDeS database file. + dry_run: If *True* (default), only report which datasets would + be removed without making any changes. Set to *False* to + actually delete the orphaned dataset records. + + Returns: + A :class:`PurgeResult` with the list of orphaned datasets and, + if *dry_run* is False, the list of datasets that were removed. + + Raises: + FileNotFoundError: If the main database file does not exist. + + """ + db_path = Path(db_path) + if not db_path.is_file(): + raise FileNotFoundError(f"Database file not found: {db_path}") + + with closing(connect(str(db_path))) as conn: + all_datasets = _build_dataset_info_list(conn) + + # Orphaned = raw_data_size_bytes is None (file not found on disk) + orphaned = [ds for ds in all_datasets if ds.raw_data_size_bytes is None] + + msg = ( + f"Found {len(all_datasets)} datasets with raw data references in {db_path}, " + f"{len(orphaned)} orphaned (file missing)." + ) + log.info(msg) + print(msg) + + removed: list[DatasetInfo] = [] + errors: list[tuple[int, Exception]] = [] + + if not dry_run and orphaned: + for ds_info in tqdm(orphaned, desc="Removing orphaned datasets"): + try: + _remove_dataset_from_db(conn, ds_info.run_id) + removed.append(ds_info) + log.debug( + "Removed orphaned dataset run_id=%d (guid=%s) from %s.", + ds_info.run_id, + ds_info.guid, + db_path, + ) + except Exception as exc: # noqa: BLE001 collect any failure and continue purging the remaining datasets + log.error("Failed to remove run_id=%d: %s", ds_info.run_id, exc) + errors.append((ds_info.run_id, exc)) + + result = PurgeResult( + total_datasets_with_raw_data=len(all_datasets), + orphaned_datasets=orphaned, + removed_datasets=removed, + dry_run=dry_run, + errors=errors, + ) + + if dry_run: + msg = f"Dry run: {len(orphaned)} orphaned datasets would be removed from {db_path}." + else: + msg = f"Removed {len(removed)} orphaned datasets from {db_path}." + log.info(msg) + print(msg) + + return result + + +def cleanup_datasets( + db_path: str | Path, + *, + older_than_days: int | None = None, + sample_name: str | None = None, + larger_than_mb: float | None = None, + dry_run: bool = True, +) -> CleanupResult: + """Remove datasets and their raw data files matching given criteria. + + This function helps manage disk space by removing datasets that match + one or more of the specified criteria. It removes both the raw data + SQLite file on disk and the corresponding records in the main database. + + Criteria are combined with AND logic: a dataset must match **all** + specified criteria to be selected for removal. Specify at least one + criterion. + + In-progress runs (started but not yet completed) are never removed, since + their raw-data file may still be actively written by another process. + + Args: + db_path: Path to the main QCoDeS database file. + older_than_days: Remove datasets whose *completed_timestamp* + (or *run_timestamp* if not completed) is older than this + many days ago. + sample_name: Remove datasets belonging to experiments with this + exact sample name. + larger_than_mb: Remove datasets whose raw data file is larger + than this many megabytes. + dry_run: If *True* (default), only report which datasets would + be removed without making any changes. Set to *False* to + actually delete datasets and their raw data files. + + Returns: + A :class:`CleanupResult` with details of the operation. + + Raises: + FileNotFoundError: If the main database file does not exist. + ValueError: If no criteria are specified, or if ``older_than_days`` or + ``larger_than_mb`` is negative. + + """ + db_path = Path(db_path) + if not db_path.is_file(): + raise FileNotFoundError(f"Database file not found: {db_path}") + + if older_than_days is None and sample_name is None and larger_than_mb is None: + raise ValueError("At least one cleanup criterion must be specified.") + + # Reject negative thresholds: a negative age puts the cutoff in the future + # and a negative size is below every file, so either would match (and, when + # dry_run=False, delete) essentially all split datasets. + if older_than_days is not None and older_than_days < 0: + raise ValueError( + f"older_than_days must be non-negative, got {older_than_days}." + ) + if larger_than_mb is not None and larger_than_mb < 0: + raise ValueError(f"larger_than_mb must be non-negative, got {larger_than_mb}.") + + with closing(connect(str(db_path))) as conn: + all_datasets = _build_dataset_info_list(conn) + + # Apply filters (AND logic) + matching: list[DatasetInfo] = [] + cutoff_ts: float | None = None + if older_than_days is not None: + cutoff_dt = datetime.now(tz=UTC) - timedelta(days=older_than_days) + cutoff_ts = cutoff_dt.timestamp() + + size_threshold_bytes: int | None = None + if larger_than_mb is not None: + size_threshold_bytes = int(larger_than_mb * 1024 * 1024) + + for ds in all_datasets: + # Never delete an in-progress run (started but not completed): + # another process may still be writing its raw-data file, and + # unlinking it would cause data loss. Such runs are skipped + # regardless of the other criteria. + if ds.completed_timestamp is None: + continue + + # Age filter + if cutoff_ts is not None: + ts = ds.completed_timestamp or ds.run_timestamp + if ts is None or ts >= cutoff_ts: + continue + + # Sample name filter + if sample_name is not None and ds.sample_name != sample_name: + continue + + # Size filter + if size_threshold_bytes is not None and ( + ds.raw_data_size_bytes is None + or ds.raw_data_size_bytes <= size_threshold_bytes + ): + continue + + matching.append(ds) + + msg = ( + f"Found {len(all_datasets)} datasets with raw data in {db_path}, " + f"{len(matching)} match cleanup criteria." + ) + log.info(msg) + print(msg) + + removed: list[DatasetInfo] = [] + errors: list[tuple[int, Exception]] = [] + total_freed: int = 0 + + if not dry_run and matching: + for ds_info in tqdm(matching, desc="Removing datasets"): + try: + raw_path = ( + Path(ds_info.raw_data_db_path) + if ds_info.raw_data_db_path + else None + ) + # Remove the DB record first: if this fails, the raw data + # file is still on disk (recoverable) rather than deleted + # while its run keeps pointing at missing data. + _remove_dataset_from_db(conn, ds_info.run_id) + # Only after the record is gone, delete the raw data file. + if raw_path is not None and raw_path.is_file(): + file_size = raw_path.stat().st_size + raw_path.unlink() + total_freed += file_size + log.debug("Deleted raw data file: %s", raw_path) + removed.append(ds_info) + log.debug( + "Removed dataset run_id=%d (guid=%s) from %s.", + ds_info.run_id, + ds_info.guid, + db_path, + ) + except Exception as exc: # noqa: BLE001 collect any failure and continue cleaning up the remaining datasets + log.error("Failed to remove run_id=%d: %s", ds_info.run_id, exc) + errors.append((ds_info.run_id, exc)) + + result = CleanupResult( + total_datasets_scanned=len(all_datasets), + matching_datasets=matching, + removed_datasets=removed, + total_size_freed_bytes=total_freed, + dry_run=dry_run, + errors=errors, + ) + + if dry_run: + total_size = sum( + ds.raw_data_size_bytes for ds in matching if ds.raw_data_size_bytes + ) + msg = f"Dry run: {len(matching)} datasets ({total_size} bytes) would be removed from {db_path}." + else: + msg = f"Removed {len(removed)} datasets from {db_path}, freed {total_freed} bytes." + log.info(msg) + print(msg) + + return result diff --git a/src/qcodes/dataset/_results_backend.py b/src/qcodes/dataset/_results_backend.py new file mode 100644 index 000000000000..c94f01248849 --- /dev/null +++ b/src/qcodes/dataset/_results_backend.py @@ -0,0 +1,330 @@ +"""Pluggable storage backends for a :class:`~qcodes.dataset.data_set.DataSet`'s +results (raw measurement data). + +A :class:`ResultsBackend` encapsulates *where and how* the results-table data of +a dataset is stored, so that :class:`~qcodes.dataset.data_set.DataSet` itself +stays free of storage-specific conditionals. The concrete backend is chosen +inside ``DataSet.__init__`` (based on config for new runs, and on the run's +recorded state for existing runs), which keeps auto-detection working for any +``DataSet(run_id=...)`` construction. + +Two backends are provided: + +- :class:`MainDatabaseResultsBackend` (the default) keeps the results table in + the main QCoDeS database. +- :class:`SeparateSqliteFileResultsBackend` writes results to an individual + per-dataset SQLite file while all metadata remains in the main database. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import TYPE_CHECKING + +from qcodes.dataset._raw_data_storage import ( + MAIN_DB_BACKEND, + PER_DATASET_DB_BACKEND, + connect_to_raw_data_db, + create_raw_data_db, + get_configured_results_backend_name, + get_raw_data_db_path, +) +from qcodes.dataset.sqlite.connection import atomic, atomic_transaction +from qcodes.dataset.sqlite.queries import ( + _check_if_table_found, + _create_run_table, + get_parameter_data, + get_raw_data_db_path_for_run, + get_shaped_parameter_data_for_one_paramtree, + set_raw_data_db_path_for_run, +) +from qcodes.dataset.sqlite.query_helpers import ( + insert_column, + insert_many_values, + length, + one, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from qcodes.dataset.data_set import DataSet + from qcodes.dataset.data_set_protocol import ParameterData + from qcodes.dataset.sqlite.connection import AtomicConnection + from qcodes.dataset.sqlite.query_helpers import VALUE + +log = logging.getLogger(__name__) + + +class ResultsBackend: + """Strategy describing where a :class:`.DataSet`'s results table lives. + + The base class keeps results in the main QCoDeS database; subclasses store + them elsewhere (e.g. a per-dataset SQLite file) while metadata stays in the + main database. A backend is owned by exactly one dataset. + + Subclasses customise only what differs: :attr:`backend_name`, + :attr:`results_conn`, the ``setup_on_*`` / :meth:`close` lifecycle, + :attr:`results_db_path`, and - if the read path differs - + :meth:`read_parameter_data`. The results-table operations below + are written against :attr:`results_conn`, so both backends reuse them + unchanged (which is why :class:`MainDatabaseResultsBackend` needs no code of + its own). + """ + + #: ``dataset.raw_data_backend`` value that selects this backend and keys its + #: settings in ``dataset.raw_data_backend_config``. + backend_name: str = MAIN_DB_BACKEND + + def __init__(self, dataset: DataSet) -> None: + self._dataset = dataset + + @property + def results_conn(self) -> AtomicConnection: + """The connection on which the results table lives.""" + return self._dataset.conn + + # -- lifecycle ------------------------------------------------------- + + def setup_on_load(self, *, read_only: bool) -> None: + """Set up the backend for an existing run being loaded. No-op here.""" + + def setup_on_new_run(self) -> None: + """Prepare the backend's storage for a newly created run. + + The default creates the (still empty) results table in the main + database so that the run is recognisable as a :class:`.DataSet` (rather + than a ``DataSetInMem``) even before it is started. Subclasses that keep + no storage in the main database override this. + """ + _create_run_table(self.results_conn, self._dataset.table_name) + + def setup_on_start(self) -> None: + """Prepare/open the backend's storage when the run is started. + + Called once all parameters are known. The default adds the parameter + columns to the table created in :meth:`setup_on_new_run`; a backend that + stores data elsewhere overrides this (e.g. to create a per-dataset file) + - and a backend that keeps no results table at all can do whatever its + storage format requires here instead. + """ + ds = self._dataset + # Add all columns in a single transaction to avoid one commit (fsync) + # per parameter. + with atomic(self.results_conn) as conn: + for spec in ds._rundescriber.interdeps.paramspecs: + insert_column(conn, ds.table_name, spec.name, spec.type) + + def close(self) -> None: + """Close any resources owned by the backend. No-op here.""" + + # -- results-table operations ---------------------------------------- + + def results_table_exists(self) -> bool: + """Whether the physical results table currently exists.""" + return _check_if_table_found(self.results_conn, self._dataset.table_name) + + def number_of_results(self) -> int: + """Number of rows in the results table (``0`` if it does not exist).""" + if not self.results_table_exists(): + return 0 + sql = f'SELECT COUNT(*) FROM "{self._dataset.table_name}"' + cursor = atomic_transaction(self.results_conn, sql) + return one(cursor, "COUNT(*)") + + def results_length(self) -> int: + """Length of the results table (max row id, ``0`` if it does not + exist).""" + if not self.results_table_exists(): + return 0 + return length(self.results_conn, self._dataset.table_name) + + def insert_results( + self, param_names: Sequence[str], values: Sequence[Sequence[VALUE]] + ) -> None: + """Insert rows of results directly (used for non-background writes).""" + insert_many_values( + self.results_conn, self._dataset.table_name, list(param_names), values + ) + + @property + def results_db_path(self) -> str | None: + """Path of the separate file holding the results table, or ``None`` when + results live in the main database. + + Lets the ``DataSet`` route background writes to the right file without + the backend knowing about the background-writer's queue-item format. + """ + return None + + def read_parameter_data( + self, + valid_param_names: list[str], + start: int | None, + end: int | None, + callback: Callable[[float], None] | None, + ) -> ParameterData: + """Read parameter data for the given parameters from the backend.""" + return get_parameter_data( + self.results_conn, + self._dataset.table_name, + valid_param_names, + start, + end, + callback, + ) + + +class MainDatabaseResultsBackend(ResultsBackend): + """Store results in the main QCoDeS database (the default). + + The reference behaviour: it adds nothing, as the base class's generic + implementation is exactly what a main-database dataset needs. + """ + + +class SeparateSqliteFileResultsBackend(ResultsBackend): + """Store results in an individual per-dataset SQLite file. + + All metadata stays in the main database; the path to the per-dataset file + is recorded in the ``raw_data_db_path`` column of the ``runs`` table (not in + the user-facing metadata), which is also how such runs are recognised when + loading. + """ + + backend_name = PER_DATASET_DB_BACKEND + + def __init__(self, dataset: DataSet) -> None: + super().__init__(dataset) + self._conn: AtomicConnection | None = None + self._db_path: str | None = None + + @property + def results_conn(self) -> AtomicConnection: + # Before the dataset is started the per-dataset file does not exist yet, + # so fall back to the main connection. The main connection has no + # results table either, which the operation methods handle via + # ``results_table_exists``. + if self._conn is not None: + return self._conn + return self._dataset.conn + + def setup_on_load(self, *, read_only: bool) -> None: + ds = self._dataset + # The path is stored in a dedicated runs-table column, not in the + # user-facing metadata. + raw_db_path = get_raw_data_db_path_for_run(ds.conn, ds.run_id) + self._db_path = raw_db_path + if raw_db_path is None: + return + if Path(raw_db_path).is_file(): + self._conn = connect_to_raw_data_db(raw_db_path, read_only=read_only) + elif ds._started: + raise FileNotFoundError( + f"Raw data file for dataset {ds.guid} not found at " + f"'{raw_db_path}'. The per-dataset SQLite file may " + f"have been moved or deleted." + ) + # else: the dataset was never started, so the raw data file has not been + # created yet - there is simply no data to connect to. + + def setup_on_new_run(self) -> None: + ds = self._dataset + # Record the raw-data backend location up front. This marks the run as a + # split-storage dataset (so it can be told apart from a DataSetInMem run, + # which also has no results table) even before it is started and before + # the raw data file is created. + raw_path_str = str(get_raw_data_db_path(ds.guid, db_path=ds.path_to_db)) + self._db_path = raw_path_str + with atomic(ds.conn) as aconn: + set_raw_data_db_path_for_run(aconn, ds.run_id, raw_path_str) + + def setup_on_start(self) -> None: + ds = self._dataset + # The raw-data path was already recorded at creation time; reuse it so + # both locations stay in sync. + raw_path_str = self._db_path or str( + get_raw_data_db_path(ds.guid, db_path=ds.path_to_db) + ) + self._conn = create_raw_data_db( + Path(raw_path_str), + ds.table_name, + ds._rundescriber.interdeps.paramspecs, + ) + if self._db_path != raw_path_str: + self._db_path = raw_path_str + with atomic(ds.conn) as aconn: + set_raw_data_db_path_for_run(aconn, ds.run_id, raw_path_str) + + @property + def results_db_path(self) -> str | None: + return self._db_path + + def read_parameter_data( + self, + valid_param_names: list[str], + start: int | None, + end: int | None, + callback: Callable[[float], None] | None, + ) -> ParameterData: + ds = self._dataset + if self._conn is None: + # Not started yet / no separate file: defer to the default reader, + # which will find no results table and return empty data. + return super().read_parameter_data(valid_param_names, start, end, callback) + # When results live in a separate DB, bypass get_parameter_data (which + # looks up the rundescriber from the main DB) and call the lower-level + # function directly with the rundescriber we already hold. + output: ParameterData = {} + for param_name in valid_param_names: + output[param_name] = get_shaped_parameter_data_for_one_paramtree( + self._conn, + ds.table_name, + ds._rundescriber, + param_name, + start, + end, + callback, + ) + return output + + def close(self) -> None: + if self._conn is not None: + self._conn.close() + + +#: Registry mapping ``dataset.raw_data_backend`` values to backend classes. +#: Add an entry here (and a matching ``raw_data_backend`` enum value plus a +#: ``raw_data_backend_config`` section in the config schema) to add a backend. +_RESULTS_BACKENDS_BY_NAME: dict[str, type[ResultsBackend]] = { + MainDatabaseResultsBackend.backend_name: MainDatabaseResultsBackend, + SeparateSqliteFileResultsBackend.backend_name: SeparateSqliteFileResultsBackend, +} + + +def select_results_backend_for_new_run(dataset: DataSet) -> ResultsBackend: + """Choose the results backend for a newly created dataset based on config.""" + name = get_configured_results_backend_name() + try: + backend_cls = _RESULTS_BACKENDS_BY_NAME[name] + except KeyError: + raise ValueError( + f"Unknown results backend {name!r} configured in " + f"'dataset.raw_data_backend'. Valid options are: " + f"{sorted(_RESULTS_BACKENDS_BY_NAME)}." + ) from None + return backend_cls(dataset) + + +def select_results_backend_for_existing_run( + dataset: DataSet, conn: AtomicConnection, run_id: int +) -> ResultsBackend: + """Choose the results backend for an existing run. + + A run whose ``raw_data_db_path`` column is set stores its results in a + separate SQLite file; otherwise the results live in the main database. + """ + if get_raw_data_db_path_for_run(conn, run_id) is not None: + return SeparateSqliteFileResultsBackend(dataset) + return MainDatabaseResultsBackend(dataset) diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py index 9364b2f85fb1..fe5339339625 100644 --- a/src/qcodes/dataset/data_set.py +++ b/src/qcodes/dataset/data_set.py @@ -60,8 +60,8 @@ get_guid_from_expid_and_counter, get_guid_from_run_id, get_metadata_from_run_id, - get_parameter_data, get_parent_dataset_links, + get_raw_data_db_path_for_run, get_run_description, get_run_timestamp_from_run_id, get_runid_from_guid, @@ -77,14 +77,20 @@ VALUE, VALUES, insert_many_values, - length, - one, select_one_where, ) from qcodes.utils import ( NumpyJSONEncoder, ) +from ._raw_data_storage import ( + connect_to_raw_data_db, +) +from ._results_backend import ( + ResultsBackend, + select_results_backend_for_existing_run, + select_results_backend_for_new_run, +) from .data_set_cache import DataSetCacheWithDBBackend from .data_set_in_memory import DataSetInMem, load_from_file from .descriptions.versioning import serialization as serial @@ -141,22 +147,42 @@ def __init__(self, queue: Queue[Any], conn: AtomicConnection): def run(self) -> None: self.conn = connect(self.path) + self._raw_data_conns: dict[str, AtomicConnection] = {} while self.keep_writing: item = self.queue.get() if item["keys"] == "stop": self.keep_writing = False self.conn.close() + for raw_conn in self._raw_data_conns.values(): + raw_conn.close() elif item["keys"] == "finalize": _WRITERS[self.path].active_datasets.remove(item["values"]) else: - self.write_results(item["keys"], item["values"], item["table_name"]) + conn = self._get_conn_for_item(item) + self.write_results( + conn, item["keys"], item["values"], item["table_name"] + ) self.queue.task_done() + def _get_conn_for_item(self, item: dict[str, Any]) -> AtomicConnection: + results_db_path = item.get("results_db_path") + if results_db_path is None: + return self.conn + if results_db_path not in self._raw_data_conns: + self._raw_data_conns[results_db_path] = connect_to_raw_data_db( + results_db_path + ) + return self._raw_data_conns[results_db_path] + def write_results( - self, keys: Sequence[str], values: Sequence[list[Any]], table_name: str + self, + conn: AtomicConnection, + keys: Sequence[str], + values: Sequence[list[Any]], + table_name: str, ) -> None: - insert_many_values(self.conn, table_name, keys, values) + insert_many_values(conn, table_name, keys, values) def shutdown(self) -> None: """ @@ -267,6 +293,10 @@ def __init__( self._debug = False self.subscribers: dict[str, _Subscriber] = {} + #: Subscriptions requested before the dataset was started, and thus + #: before the results table exists. They are materialised into real + #: subscribers (with their SQL triggers) once the dataset is started. + self._pending_subscribers: dict[str, dict[str, Any]] = {} self._parent_dataset_links: list[Link] #: In memory representation of the data in the dataset. self._cache: DataSetCacheWithDBBackend = DataSetCacheWithDBBackend(self) @@ -290,6 +320,13 @@ def __init__( self._export_info = ExportInfo.from_str( self.metadata.get("export_info", "") ) + # Select the results backend from the run's recorded state, so that + # e.g. a run whose data lives in a separate SQLite file is detected + # automatically here (no matter how the DataSet is constructed). + self._results_backend: ResultsBackend = ( + select_results_backend_for_existing_run(self, self.conn, self.run_id) + ) + self._results_backend.setup_on_load(read_only=read_only) else: # Actually perform all the side effects needed for the creation # of a new dataset. Note that a dataset is created (in the DB) @@ -298,6 +335,11 @@ def __init__( if exp_id is None: exp_id = get_default_experiment_id(self.conn) name = name or "dataset" + # Select the results backend from config. The backend sets up its + # own storage when the run is started (see ``setup_on_start``), not + # here - ``create_run`` only records the run metadata, mirroring how + # ``DataSetInMem`` records runs. + self._results_backend = select_results_backend_for_new_run(self) _, run_id, __ = create_run( self.conn, exp_id, @@ -306,6 +348,7 @@ def __init__( parameters=None, values=values, metadata=metadata, + create_run_table=False, ) # this is really the UUID (an ever increasing count in the db) self._run_id = run_id @@ -324,6 +367,10 @@ def __init__( self._metadata = get_metadata_from_run_id(self.conn, self.run_id) self._parent_dataset_links = [] self._export_info = ExportInfo({}) + + # Let the backend record any bookkeeping for the new run (e.g. the + # location of the per-dataset raw data file). + self._results_backend.setup_on_new_run() assert self.path_to_db is not None if _WRITERS.get(self.path_to_db) is None: queue: Queue[Any] = Queue() @@ -358,6 +405,28 @@ def prepare( def cache(self) -> DataSetCacheWithDBBackend: return self._cache + @property + def _results_conn(self) -> AtomicConnection: + """The connection on which this dataset's results table lives. + + Delegates to the results backend: the main database connection by + default, or a separate per-dataset connection when the results are + stored elsewhere. Collaborators that operate on the results table + directly (the cache and subscribers) use this. + """ + return self._results_backend.results_conn + + @property + def _results_table_exists(self) -> bool: + """Whether the physical results table currently exists. + + With a separate results backend the table is created only once the + dataset has been started, so before that (and in the main database in + general) no results table exists. Callers that count rows must handle + this case. + """ + return self._results_backend.results_table_exists() + @property def run_id(self) -> int: return self._run_id @@ -419,9 +488,7 @@ def snapshot_raw(self) -> str | None: @property def number_of_results(self) -> int: - sql = f'SELECT COUNT(*) FROM "{self.table_name}"' - cursor = atomic_transaction(self.conn, sql) - return one(cursor, "COUNT(*)") + return self._results_backend.number_of_results() @property def counter(self) -> int: @@ -683,10 +750,24 @@ def _perform_start_actions(self, start_bg_writer: bool) -> None: """ paramspecs = new_to_old(self._rundescriber.interdeps).paramspecs - for spec in paramspecs: - add_parameter( - spec, conn=self.conn, run_id=self.run_id, insert_into_results_table=True - ) + # Register all parameters in a single transaction (each add_parameter + # would otherwise commit on its own) to avoid one fsync per parameter. + with atomic(self.conn) as conn: + for spec in paramspecs: + add_parameter( + spec, + conn=conn, + run_id=self.run_id, + # The results table is created wholesale by the backend + # just below, so only the layouts/dependencies are updated + # here. + insert_into_results_table=False, + ) + + # Let the backend set up its storage (add columns to the main-database + # results table, create a per-dataset SQLite file, ...) now that all + # parameters are known. + self._results_backend.setup_on_start() desc_str = serial.to_json_for_storage(self.description) @@ -723,6 +804,11 @@ def _perform_start_actions(self, start_bg_writer: bool) -> None: writer_status.active_datasets.add(self.run_id) self.cache.prepare() + # Now that the dataset is started (and, for split raw data storage, the + # results table exists in the per-dataset file), create any subscribers + # that were requested while the dataset was still pristine. + self._start_pending_subscribers() + def mark_completed(self) -> None: """ Mark :class:`.DataSet` as complete and thus read-only and notify the subscribers @@ -770,14 +856,17 @@ def add_results(self, results: Sequence[Mapping[str, VALUE]]) -> None: writer_status = self._writer_status if writer_status.write_in_background: - item = { + item: dict[str, Any] = { "keys": list(expected_keys), "values": values, "table_name": self.table_name, + # None means the main database; a separate backend gives the + # per-dataset file the background writer should write to. + "results_db_path": self._results_backend.results_db_path, } writer_status.data_write_queue.put(item) else: - insert_many_values(self.conn, self.table_name, list(expected_keys), values) + self._results_backend.insert_results(list(expected_keys), values) def _raise_if_not_writable(self) -> None: if self.pristine: @@ -869,8 +958,9 @@ def get_parameter_data( else: valid_param_names = self._validate_parameters(*params) - return get_parameter_data( - self.conn, self.table_name, valid_param_names, start, end, callback + + return self._results_backend.read_parameter_data( + valid_param_names, start, end, callback ) def to_pandas_dataframe_dict( @@ -1151,12 +1241,70 @@ def subscribe( callback_kwargs: Mapping[str, Any] | None = None, ) -> str: subscriber_id = uuid.uuid4().hex + if not self._results_table_exists: + # The results table does not exist yet - this happens with split + # raw data storage, where the table lives in the per-dataset file + # created only when the dataset is started. A subscriber has + # nothing to observe before the dataset is started, so defer + # creating it (and its SQL trigger) until then. + self._queue_pending_subscriber( + subscriber_id, callback, min_wait, min_count, state, callback_kwargs + ) + return subscriber_id + self._create_and_start_subscriber( + subscriber_id, callback, min_wait, min_count, state, callback_kwargs + ) + return subscriber_id + + def _create_and_start_subscriber( + self, + subscriber_id: str, + callback: Callable[[Any, int, Any | None], None], + min_wait: int, + min_count: int, + state: Any | None, + callback_kwargs: Mapping[str, Any] | None, + ) -> None: + """Create a :class:`_Subscriber` (and its SQL trigger) and start it.""" subscriber = _Subscriber( self, subscriber_id, callback, state, min_wait, min_count, callback_kwargs ) self.subscribers[subscriber_id] = subscriber subscriber.start() - return subscriber_id + + def _queue_pending_subscriber( + self, + subscriber_id: str, + callback: Callable[[Any, int, Any | None], None], + min_wait: int, + min_count: int, + state: Any | None, + callback_kwargs: Mapping[str, Any] | None, + ) -> None: + """Record a subscription requested before the dataset was started, to + be materialised once the results table exists (see + :meth:`_start_pending_subscribers`).""" + self._pending_subscribers[subscriber_id] = { + "callback": callback, + "min_wait": min_wait, + "min_count": min_count, + "state": state, + "callback_kwargs": callback_kwargs, + } + + def _start_pending_subscribers(self) -> None: + """Materialise subscriptions that were requested before the dataset was + started (and thus before the results table existed).""" + for subscriber_id, kwargs in self._pending_subscribers.items(): + self._create_and_start_subscriber( + subscriber_id, + kwargs["callback"], + kwargs["min_wait"], + kwargs["min_count"], + kwargs["state"], + kwargs["callback_kwargs"], + ) + self._pending_subscribers.clear() def subscribe_from_config(self, name: str) -> str: """ @@ -1196,7 +1344,11 @@ def unsubscribe(self, uuid: str) -> None: """ Remove subscriber with the provided uuid """ - with atomic(self.conn) as conn: + if uuid in self._pending_subscribers: + # Not yet materialised into a real subscriber/trigger. + del self._pending_subscribers[uuid] + return + with atomic(self._results_conn) as conn: sub = self.subscribers[uuid] remove_trigger(conn, sub.trigger_id) sub.schedule_stop() @@ -1207,12 +1359,16 @@ def unsubscribe_all(self) -> None: """ Remove all subscribers """ + # Drop subscriptions that were deferred before the dataset started and + # never materialised into real subscribers/triggers. + self._pending_subscribers.clear() sql = """ SELECT name FROM sqlite_master WHERE type = 'trigger' """ - triggers = atomic_transaction(self.conn, sql).fetchall() - with atomic(self.conn) as conn: + data_conn = self._results_conn + triggers = atomic_transaction(data_conn, sql).fetchall() + with atomic(data_conn) as conn: for (trigger,) in triggers: remove_trigger(conn, trigger) for sub in self.subscribers.values(): @@ -1225,7 +1381,7 @@ def get_metadata(self, tag: str) -> VALUE | None: return get_data_by_tag_and_table_name(self.conn, tag, self.table_name) def __len__(self) -> int: - return length(self.conn, self.table_name) + return self._results_backend.results_length() def __repr__(self) -> str: out = [] @@ -1633,7 +1789,7 @@ def load_by_run_spec( ) if len(guids) == 1: - d = load_by_guid(guids[0], internal_conn) + d = load_by_guid(guids[0], internal_conn, read_only=read_only) elif len(guids) > 1: print(generate_dataset_table(guids, conn=internal_conn)) raise NameError( @@ -1746,7 +1902,7 @@ def load_by_id( raise ValueError( f"Run with run_id {run_id} does not exist in the database: {internal_conn.path_to_dbfile}" ) - d = _get_datasetprotocol_from_guid(guid, internal_conn) + d = _get_datasetprotocol_from_guid(guid, internal_conn, read_only=read_only) finally: # dataset takes ownership of the connection but DataSetInMem does not if not conn and not isinstance(d, DataSet): @@ -1789,7 +1945,7 @@ def load_by_guid( # this function raises a RuntimeError if more than one run matches the GUID try: - d = _get_datasetprotocol_from_guid(guid, internal_conn) + d = _get_datasetprotocol_from_guid(guid, internal_conn, read_only=read_only) finally: # dataset takes ownership of the connection but DataSetInMem does not if not conn and not isinstance(d, DataSet): @@ -1839,7 +1995,7 @@ def load_by_counter( # this function raises a RuntimeError if more than one run matches the GUID try: guid = get_guid_from_expid_and_counter(internal_conn, exp_id, counter) - d = _get_datasetprotocol_from_guid(guid, internal_conn) + d = _get_datasetprotocol_from_guid(guid, internal_conn, read_only=read_only) finally: # dataset takes ownership of the connection but DataSetInMem does not if not conn and not isinstance(d, DataSet): @@ -1849,7 +2005,7 @@ def load_by_counter( def _get_datasetprotocol_from_guid( - guid: str, conn: AtomicConnection + guid: str, conn: AtomicConnection, *, read_only: bool = False ) -> DataSetProtocol: run_id = get_runid_from_guid(conn, guid) if run_id is None: @@ -1875,8 +2031,15 @@ def _get_datasetprotocol_from_guid( return d result_table_name = _get_result_table_name_by_guid(conn, guid) - if _check_if_table_found(conn, result_table_name): - d = DataSet(conn=conn, run_id=run_id) + # The results table is absent from the main DB when raw data is stored in a + # separate per-dataset SQLite file. Such runs are marked with a + # raw_data_db_path; anything else without a results table is an in-memory + # (netcdf-backed) dataset. + if ( + _check_if_table_found(conn, result_table_name) + or get_raw_data_db_path_for_run(conn, run_id) is not None + ): + d = DataSet(conn=conn, run_id=run_id, read_only=read_only) else: d = DataSetInMem._load_from_db(conn=conn, guid=guid) diff --git a/src/qcodes/dataset/data_set_cache.py b/src/qcodes/dataset/data_set_cache.py index 3456909d2f0c..3d8f506b2c38 100644 --- a/src/qcodes/dataset/data_set_cache.py +++ b/src/qcodes/dataset/data_set_cache.py @@ -552,12 +552,15 @@ def load_data_from_db(self) -> None: self._loaded_from_completed_ds = True if self._data == {}: self.prepare() + # The connection on which the dataset's results table lives (the main + # database, or a separate per-dataset file). + data_conn = self._dataset._results_conn ( self._write_status, self._read_status, self._data, ) = load_new_data_from_db_and_append( - self._dataset.conn, + data_conn, self._dataset.table_name, self.rundescriber, self._write_status, diff --git a/src/qcodes/dataset/database_extract_runs.py b/src/qcodes/dataset/database_extract_runs.py index 0c25bf58af86..5b37740815bb 100644 --- a/src/qcodes/dataset/database_extract_runs.py +++ b/src/qcodes/dataset/database_extract_runs.py @@ -166,8 +166,6 @@ def _extract_single_dataset_into_db( f"GUID: {dataset.guid} and run_id: {dataset.run_id}" ) - source_conn = dataset.conn - run_id = get_runid_from_guid(target_conn, dataset.guid) if run_id is not None: @@ -178,7 +176,7 @@ def _extract_single_dataset_into_db( ) assert target_table_name is not None _populate_results_table( - source_conn, target_conn, dataset.table_name, target_table_name + dataset._results_conn, target_conn, dataset.table_name, target_table_name ) diff --git a/src/qcodes/dataset/export_config.py b/src/qcodes/dataset/export_config.py index 36dba5f38e31..c8ba2817d4ea 100644 --- a/src/qcodes/dataset/export_config.py +++ b/src/qcodes/dataset/export_config.py @@ -91,10 +91,17 @@ def get_data_export_automatic() -> bool: return export_automatic -def _expand_export_path(export_path: str) -> str: - db_location = Path(qcodes.config["core"]["db_location"]).expanduser().absolute() - expanded_export_folder = db_location.parent / "_".join( - (db_location.stem, db_location.suffix.replace(".", "")) +def _expand_export_path(export_path: str, db_location: str | None = None) -> str: + # db_location defaults to the global config, but callers can pass the + # location of a specific database (e.g. a dataset's own file) so that + # {db_location} resolves relative to that database rather than the global + # config value. + location: str = ( + db_location if db_location is not None else qcodes.config["core"]["db_location"] + ) + db_file = Path(location).expanduser().absolute() + expanded_export_folder = db_file.parent / "_".join( + (db_file.stem, db_file.suffix.replace(".", "")) ) # we ignore the rule here since this is actually a manual expansion of the f string return export_path.replace("{db_location}", str(expanded_export_folder)) # noqa: RUF027 diff --git a/src/qcodes/dataset/sqlite/database.py b/src/qcodes/dataset/sqlite/database.py index 86e6405e0511..cb39c378c09d 100644 --- a/src/qcodes/dataset/sqlite/database.py +++ b/src/qcodes/dataset/sqlite/database.py @@ -120,19 +120,47 @@ def _adapt_complex(value: complex | np.complexfloating) -> sqlite3.Binary: return sqlite3.Binary(out.read()) -def connect( - name: str | Path, debug: bool = False, version: int = -1, read_only: bool = False +def _register_numpy_sqlite_adapters_and_converters() -> None: + """ + Register the numpy/sqlite type adapters and converters that QCoDeS needs + to round-trip array, numeric and complex data. The registrations are + global to the ``sqlite3`` module and idempotent, so calling this multiple + times is safe. + """ + # register numpy->binary(TEXT) adapter + sqlite3.register_adapter(np.ndarray, _adapt_array) + # register binary(TEXT) -> numpy converter + sqlite3.register_converter("array", _convert_array) + + # Make sure numpy ints and floats types are inserted properly + for numpy_int in numpy_ints: + sqlite3.register_adapter(numpy_int, int) + + sqlite3.register_converter("numeric", _convert_numeric) + + for numpy_float in (float, *numpy_floats): + sqlite3.register_adapter(numpy_float, _adapt_float) + + for complex_type in complex_types: + # https://github.com/python/typeshed/issues/2429 + sqlite3.register_adapter(complex_type, _adapt_complex) # type: ignore[arg-type] + sqlite3.register_converter("complex", _convert_complex) + + +def _connect_to_sqlite_file( + name: str | Path, read_only: bool = False ) -> AtomicConnection: """ - Connect or create database. If debug the queries will be echoed back. - This function takes care of registering the numpy/sqlite type - converters that we need. + Open a bare :class:`AtomicConnection` to an sqlite file with the settings + that QCoDeS relies on (URI mode, ``PARSE_DECLTYPES``) and the numpy type + adapters/converters registered. + + Unlike :func:`connect`, this does **not** create or upgrade any QCoDeS + schema; it is the low-level building block shared by :func:`connect` and + the per-dataset raw-data connection helper. Args: name: name or path to the sqlite file - debug: should tracing be turned on. - version: which version to create. We count from 0. -1 means 'latest'. - Should always be left at -1 except when testing. read_only: Should the database be opened in read only mode. Returns: @@ -140,17 +168,14 @@ def connect( :class:`AtomicConnection`, which is a subclass of :class:`sqlite3.Connection`) """ - # register numpy->binary(TEXT) adapter - sqlite3.register_adapter(np.ndarray, _adapt_array) - # register binary(TEXT) -> numpy converter - sqlite3.register_converter("array", _convert_array) + _register_numpy_sqlite_adapters_and_converters() path = f"file:{name!s}" if read_only: path = path + "?mode=ro" - conn = sqlite3.connect( + return sqlite3.connect( path, detect_types=sqlite3.PARSE_DECLTYPES, check_same_thread=True, @@ -158,6 +183,29 @@ def connect( factory=AtomicConnection, ) + +def connect( + name: str | Path, debug: bool = False, version: int = -1, read_only: bool = False +) -> AtomicConnection: + """ + Connect or create database. If debug the queries will be echoed back. + This function takes care of registering the numpy/sqlite type + converters that we need. + + Args: + name: name or path to the sqlite file + debug: should tracing be turned on. + version: which version to create. We count from 0. -1 means 'latest'. + Should always be left at -1 except when testing. + read_only: Should the database be opened in read only mode. + + Returns: + connection object to the database (note, it is + :class:`AtomicConnection`, which is a subclass of :class:`sqlite3.Connection`) + + """ + conn = _connect_to_sqlite_file(name, read_only=read_only) + latest_supported_version = _latest_available_version() db_version = get_user_version(conn) @@ -168,20 +216,6 @@ def connect( f"version {latest_supported_version}" ) - # Make sure numpy ints and floats types are inserted properly - for numpy_int in numpy_ints: - sqlite3.register_adapter(numpy_int, int) - - sqlite3.register_converter("numeric", _convert_numeric) - - for numpy_float in (float, *numpy_floats): - sqlite3.register_adapter(numpy_float, _adapt_float) - - for complex_type in complex_types: - # https://github.com/python/typeshed/issues/2429 - sqlite3.register_adapter(complex_type, _adapt_complex) # type: ignore[arg-type] - sqlite3.register_converter("complex", _convert_complex) - if debug: conn.set_trace_callback(print) diff --git a/src/qcodes/dataset/sqlite/db_overview.py b/src/qcodes/dataset/sqlite/db_overview.py index e596439f46aa..088bdae9ed3f 100644 --- a/src/qcodes/dataset/sqlite/db_overview.py +++ b/src/qcodes/dataset/sqlite/db_overview.py @@ -18,19 +18,23 @@ import logging import sqlite3 from contextlib import closing, nullcontext +from pathlib import Path from typing import TYPE_CHECKING from typing_extensions import TypedDict -from qcodes.dataset.sqlite.database import conn_from_dbpath_or_conn +from qcodes.dataset.sqlite.database import ( + _connect_to_sqlite_file, + conn_from_dbpath_or_conn, +) from qcodes.dataset.sqlite.query_helpers import is_column_in_table if TYPE_CHECKING: from collections.abc import Sequence - from pathlib import Path from qcodes.dataset.sqlite.connection import AtomicConnection + log = logging.getLogger(__name__) @@ -179,6 +183,13 @@ def get_db_overview( ] extra_select = "".join(f", r.{col}" for col in valid_extra_columns) + # Split-storage runs keep their results table in a per-dataset SQLite + # file recorded in ``raw_data_db_path`` (the column only exists once the + # feature has been used). Select it last so the fixed/extra column + # indices below are unaffected. + has_raw_col = is_column_in_table(c, "runs", "raw_data_db_path") + raw_select = ", r.raw_data_db_path" if has_raw_col else "" + # ``run_description`` is queried to derive the record count for # completed runs; the (potentially large) snapshot is deliberately # excluded. @@ -186,7 +197,7 @@ def get_db_overview( SELECT r.run_id, e.name, e.sample_name, r.name, r.run_timestamp, r.completed_timestamp, r.guid, r.result_table_name, - r.run_description{extra_select} + r.run_description{extra_select}{raw_select} FROM runs r JOIN experiments e ON r.exp_id = e.exp_id WHERE r.run_id > ? @@ -199,26 +210,48 @@ def get_db_overview( log.warning("Could not query database overview: %s", e) return overview + n_fixed = 9 # number of columns selected before ``extra_columns`` + raw_idx = n_fixed + len(valid_extra_columns) if has_raw_col else None + # ``result_counter`` in the runs table is the run's ordinal within its # experiment, not a data-point count, so it is not usable here. For the # ``array`` paramtype a single INSERT can also contain many data points. # The real number of data points is therefore the row count of the - # results table, queried separately. - result_tables = {row[7] for row in rows if row[7]} - row_counts: dict[str, int] = {} - for table in result_tables: - try: - (count,) = c.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone() - except sqlite3.Error: - continue # results table may not exist (yet) - row_counts[table] = count + # results table, queried separately (keyed by run_id). For split-storage + # runs the table lives in the recorded per-dataset file, so it is + # counted there rather than on the main connection. + row_counts: dict[int, int] = {} + raw_conns: dict[str, AtomicConnection] = {} + try: + for row in rows: + table = row[7] + if not table: + continue + raw_db_path = row[raw_idx] if raw_idx is not None else None + count_conn = c + if raw_db_path: + if not Path(raw_db_path).is_file(): + continue # per-dataset file moved/deleted + if raw_db_path not in raw_conns: + raw_conns[raw_db_path] = _connect_to_sqlite_file( + raw_db_path, read_only=True + ) + count_conn = raw_conns[raw_db_path] + try: + (count,) = count_conn.execute( + f'SELECT COUNT(*) FROM "{table}"' + ).fetchone() + except sqlite3.Error: + continue # results table may not exist (yet) + row_counts[row[0]] = count + finally: + for raw_conn in raw_conns.values(): + raw_conn.close() - n_fixed = 9 # number of columns selected before ``extra_columns`` for row in rows: run_id = row[0] started_date, started_time = _format_timestamp(row[4]) completed_date, completed_time = _format_timestamp(row[5]) - result_table = row[7] or "" is_completed = row[5] is not None and row[5] != 0 # The record count is a best-effort data-point count. For completed @@ -228,9 +261,9 @@ def get_db_overview( if is_completed: records = _records_from_run_description(row[8]) if records == 0: - records = row_counts.get(result_table, 0) + records = row_counts.get(run_id, 0) else: - records = row_counts.get(result_table, 0) + records = row_counts.get(run_id, 0) if records == 0: records = _records_from_run_description(row[8]) diff --git a/src/qcodes/dataset/sqlite/queries.py b/src/qcodes/dataset/sqlite/queries.py index c42eeb676f17..ac7a63618ac8 100644 --- a/src/qcodes/dataset/sqlite/queries.py +++ b/src/qcodes/dataset/sqlite/queries.py @@ -12,7 +12,7 @@ import unicodedata import warnings from itertools import zip_longest -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast import numpy as np import numpy.typing as npt @@ -90,6 +90,12 @@ "captured_counter", ) +#: Name of the ``runs``-table column used to record the location of a dataset's +#: raw data when it is stored outside the main database (e.g. in a per-dataset +#: SQLite file). It is an internal storage detail and is deliberately excluded +#: from the user-facing dataset metadata. +RAW_DATA_DB_PATH_COLUMN = "raw_data_db_path" + def is_run_id_in_database(conn: AtomicConnection, *run_ids: int) -> dict[int, bool]: """ @@ -1838,7 +1844,7 @@ def get_metadata_from_run_id(conn: AtomicConnection, run_id: int) -> dict[str, A """ Get all metadata associated with the specified run """ - non_metadata = RUNS_TABLE_COLUMNS + non_metadata = (*RUNS_TABLE_COLUMNS, RAW_DATA_DB_PATH_COLUMN) metadata = {} possible_tags = [] @@ -2328,3 +2334,130 @@ def _get_result_table_name_by_guid(conn: AtomicConnection, guid: str) -> str: sql = "SELECT result_table_name FROM runs WHERE guid=?" formatted_name = one(transaction(conn, sql, guid), "result_table_name") return formatted_name + + +def get_raw_data_db_path_for_run(conn: AtomicConnection, run_id: int) -> str | None: + """Return the stored raw-data file path for a run. + + The path is read directly from the ``raw_data_db_path`` column of the + ``runs`` table (which is kept out of the user-facing metadata). Returns + ``None`` if the column does not exist or is not set for the run. + """ + if not is_column_in_table(conn, "runs", RAW_DATA_DB_PATH_COLUMN): + return None + cursor = conn.execute( + f'SELECT "{RAW_DATA_DB_PATH_COLUMN}" FROM runs WHERE run_id = ?', (run_id,) + ) + row = cursor.fetchone() + if row is None: + return None + return row[0] + + +def set_raw_data_db_path_for_run( + conn: AtomicConnection, run_id: int, raw_data_db_path: str +) -> None: + """Record the raw-data file path for a run in the ``runs`` table.""" + add_data_to_dynamic_columns( + conn, run_id, {RAW_DATA_DB_PATH_COLUMN: raw_data_db_path} + ) + + +class RawDataDatasetRecord(NamedTuple): + """A run in the main database whose raw data is stored in a separate file. + + Returned by :func:`get_datasets_with_raw_data_path` for each started run + that records a ``raw_data_db_path``. + """ + + run_id: int + guid: str + experiment_name: str + sample_name: str + run_timestamp: float | None + completed_timestamp: float | None + result_table_name: str + raw_data_db_path: str + + +def get_datasets_with_raw_data_path( + conn: AtomicConnection, +) -> list[RawDataDatasetRecord]: + """Get all datasets that have a raw_data_db_path metadata column set. + + Only datasets that have been started (``run_timestamp`` is set) are + returned, since an unstarted run records the intended raw-data path but + never creates the corresponding file. + + Returns: + A list of :class:`RawDataDatasetRecord`. Returns an empty list if the + ``raw_data_db_path`` column does not exist. + + """ + if not is_column_in_table(conn, "runs", "raw_data_db_path"): + return [] + + sql = """ + SELECT r.run_id, r.guid, e.name, e.sample_name, + r.run_timestamp, r.completed_timestamp, + r.result_table_name, r.raw_data_db_path + FROM runs r + JOIN experiments e ON r.exp_id = e.exp_id + WHERE r.raw_data_db_path IS NOT NULL + AND r.run_timestamp IS NOT NULL + """ + cursor = atomic_transaction(conn, sql) + return [RawDataDatasetRecord(*row) for row in cursor.fetchall()] + + +def _remove_dataset_from_db(conn: AtomicConnection, run_id: int) -> None: + """Remove a single dataset's records from the database. + + Internal helper for the raw-data management functions + (:func:`~qcodes.dataset.purge_orphaned_datasets`, + :func:`~qcodes.dataset.cleanup_datasets`); not part of the public API. + + Deletes the run row, associated layouts and dependencies, and drops + the results table (if it exists). The results-table name is looked up + from the ``runs`` table. + + Args: + conn: Connection to the database. + run_id: The run_id of the dataset to remove. + + """ + result_table_name = select_one_where( + conn, "runs", "result_table_name", "run_id", run_id + ) + assert isinstance(result_table_name, str) + + with atomic(conn) as aconn: + # Guard against dropping an unintended table if result_table_name is + # malformed, reusing the same validation used when creating tables. + _validate_table_name(result_table_name) + + # Get layout_ids for this run (needed for dependencies) + cursor = transaction( + aconn, "SELECT layout_id FROM layouts WHERE run_id = ?", run_id + ) + layout_ids = [row[0] for row in cursor.fetchall()] + + # Delete dependencies referencing these layouts + if layout_ids: + placeholders = ",".join("?" * len(layout_ids)) + transaction( + aconn, + f"DELETE FROM dependencies WHERE dependent IN ({placeholders})" + f" OR independent IN ({placeholders})", + *layout_ids, + *layout_ids, + ) + + # Delete layouts + transaction(aconn, "DELETE FROM layouts WHERE run_id = ?", run_id) + + # Drop the results table in the DB (if it exists) + transaction(aconn, f'DROP TABLE IF EXISTS "{result_table_name}"') + + # Delete the run row + transaction(aconn, "DELETE FROM runs WHERE run_id = ?", run_id) diff --git a/src/qcodes/dataset/subscriber.py b/src/qcodes/dataset/subscriber.py index 6b540599396c..5c87049c3f65 100644 --- a/src/qcodes/dataset/subscriber.py +++ b/src/qcodes/dataset/subscriber.py @@ -64,7 +64,7 @@ def __init__( self.callback_id = f"callback{self._id}" self.trigger_id = f"sub{self._id}" - conn = dataSet.conn + conn = dataSet._results_conn conn.create_function(self.callback_id, -1, self._cache_data_to_queue) diff --git a/tests/dataset/test_raw_data_storage.py b/tests/dataset/test_raw_data_storage.py new file mode 100644 index 000000000000..28382dec6dd5 --- /dev/null +++ b/tests/dataset/test_raw_data_storage.py @@ -0,0 +1,1028 @@ +""" +Tests for per-dataset raw data SQLite storage. + +When ``dataset.raw_data_backend`` is set to ``"sqlite_per_dataset_db"``, +measurement data (results tables) are written to individual SQLite files while +metadata remains in the main database. +""" + +from __future__ import annotations + +import gc +import shutil +import sqlite3 +import time +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +import qcodes as qc +from qcodes.dataset import new_data_set, new_experiment +from qcodes.dataset._raw_data_storage import ( + cleanup_datasets, + connect_to_raw_data_db, + create_raw_data_db, + get_configured_results_backend_name, + get_raw_data_db_path, + get_raw_data_folder, + get_results_backend_config, + is_raw_data_storage_enabled, + purge_orphaned_datasets, + update_raw_data_paths, +) +from qcodes.dataset._results_backend import ( + MainDatabaseResultsBackend, + SeparateSqliteFileResultsBackend, +) +from qcodes.dataset.data_set import DataSet, load_by_id +from qcodes.dataset.database_extract_runs import ( + export_datasets_and_create_metadata_db, + extract_runs_into_db, +) +from qcodes.dataset.descriptions.dependencies import InterDependencies_ +from qcodes.dataset.sqlite.database import connect, initialise_database +from qcodes.dataset.sqlite.db_overview import get_db_overview +from qcodes.dataset.sqlite.queries import get_raw_data_db_path_for_run +from qcodes.dataset.sqlite.query_helpers import select_one_where +from qcodes.parameters import ParamSpecBase + +if TYPE_CHECKING: + from collections.abc import Generator + + +def _raw_file(ds: DataSet) -> Path: + """Return the per-dataset raw data file path, asserting it is set.""" + raw_path = get_raw_data_db_path_for_run(ds.conn, ds.run_id) + assert raw_path is not None + return Path(raw_path) + + +def _uses_separate_file(ds: DataSet) -> bool: + """Whether the dataset stores results in a separate per-dataset file.""" + return isinstance(ds._results_backend, SeparateSqliteFileResultsBackend) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def _raw_data_db(tmp_path: Path) -> Generator[None, None, None]: + """Set up a temp DB with the per-dataset SQLite results backend selected.""" + db_path = str(tmp_path / "test.db") + qc.config["core"]["db_location"] = db_path + qc.config["core"]["db_debug"] = False + qc.config["dataset"]["raw_data_backend"] = "sqlite_per_dataset_db" + qc.config["dataset"]["raw_data_backend_config"] = { + "sqlite_per_dataset_db": {"raw_data_path": str(tmp_path / "raw_data")} + } + initialise_database() + try: + yield + finally: + qc.config["dataset"]["raw_data_backend"] = "sqlite_main_db" + qc.config["dataset"]["raw_data_backend_config"] = { + "sqlite_per_dataset_db": {"raw_data_path": "{db_location}"} + } + gc.collect() + + +@pytest.fixture() +def _raw_data_experiment(_raw_data_db: None) -> Generator[None, None, None]: + """Create a test experiment inside the raw data DB.""" + e = new_experiment("test-experiment", sample_name="test-sample") + try: + yield + finally: + e.conn.close() + + +# --------------------------------------------------------------------------- +# Unit tests - raw_data_storage module +# --------------------------------------------------------------------------- + + +class TestRawDataStorageHelpers: + def test_is_raw_data_storage_enabled_default(self, tmp_path: Path) -> None: + assert not is_raw_data_storage_enabled() + + def test_is_raw_data_storage_enabled_on(self, _raw_data_db: None) -> None: + assert is_raw_data_storage_enabled() + + def test_configured_backend_default(self, tmp_path: Path) -> None: + assert get_configured_results_backend_name() == "sqlite_main_db" + + def test_configured_backend_selected(self, _raw_data_db: None) -> None: + assert get_configured_results_backend_name() == "sqlite_per_dataset_db" + + def test_backend_config_holds_raw_data_path( + self, _raw_data_db: None, tmp_path: Path + ) -> None: + config = get_results_backend_config("sqlite_per_dataset_db") + assert config["raw_data_path"] == str(tmp_path / "raw_data") + # A backend without configured settings gets an empty mapping. + assert get_results_backend_config("sqlite_main_db") == {} + + def test_unknown_backend_raises(self, tmp_path: Path) -> None: + qc.config["core"]["db_location"] = str(tmp_path / "unknown.db") + qc.config["core"]["db_debug"] = False + initialise_database() + new_experiment("e", sample_name="s") + qc.config["dataset"]["raw_data_backend"] = "does_not_exist" + try: + with pytest.raises(ValueError, match="Unknown results backend"): + new_data_set("ds") + finally: + qc.config["dataset"]["raw_data_backend"] = "sqlite_main_db" + gc.collect() + + def test_get_raw_data_folder(self, _raw_data_db: None, tmp_path: Path) -> None: + folder = get_raw_data_folder() + assert folder == (tmp_path / "raw_data") + + def test_get_raw_data_db_path(self, tmp_path: Path) -> None: + folder = tmp_path / "raw_data" + path = get_raw_data_db_path("abc-123", folder) + assert path == folder / "abc-123.db" + + def test_raw_data_folder_derives_from_given_db_path(self, tmp_path: Path) -> None: + """With the default ``{db_location}`` template, the folder resolves next + to the owning dataset's database, not the global config location.""" + ds_db = tmp_path / "sub" / "exp.db" + assert get_raw_data_folder(str(ds_db)) == tmp_path / "sub" / "exp_db" + assert ( + get_raw_data_db_path("g", db_path=str(ds_db)) + == tmp_path / "sub" / "exp_db" / "g.db" + ) + + def test_create_raw_data_db(self, tmp_path: Path) -> None: + db_path = tmp_path / "raw" / "test.db" + params = [ + ParamSpecBase("x", "numeric"), + ParamSpecBase("y", "numeric"), + ] + conn = create_raw_data_db(db_path, "results_table", params) + try: + # Verify table exists and has the right columns + cursor = conn.execute("PRAGMA table_info('results_table')") + columns = {row[1] for row in cursor.fetchall()} + assert "id" in columns + assert "x" in columns + assert "y" in columns + finally: + conn.close() + + def test_connect_to_raw_data_db(self, tmp_path: Path) -> None: + db_path = tmp_path / "test.db" + conn = connect_to_raw_data_db(db_path) + try: + # Should be a valid connection with numpy adapters + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, val REAL)") + conn.execute("INSERT INTO t (val) VALUES (?)", (3.14,)) + conn.commit() + cursor = conn.execute("SELECT val FROM t") + assert cursor.fetchone()[0] == pytest.approx(3.14) + finally: + conn.close() + + def test_connect_to_raw_data_db_read_only(self, tmp_path: Path) -> None: + db_path = tmp_path / "test.db" + # Create first + conn = connect_to_raw_data_db(db_path) + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY)") + conn.commit() + conn.close() + # Now open read-only + conn = connect_to_raw_data_db(db_path, read_only=True) + try: + with pytest.raises(sqlite3.OperationalError): + conn.execute("INSERT INTO t (id) VALUES (1)") + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Integration tests - DataSet with split raw data +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_raw_data_experiment") +class TestDataSetWithSplitRawData: + def _make_dataset_with_data( + self, n_rows: int = 10 + ) -> tuple[DataSet, list[dict[str, float]]]: + """Create a dataset, add data, mark completed, return ds and data.""" + ds = new_data_set("test-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + + results = [{"x": float(i), "y": float(i**2)} for i in range(n_rows)] + ds.add_results(results) + ds.mark_completed() + return ds, results + + @staticmethod + def _close_ds(ds: DataSet) -> None: + """Close both main and raw data connections.""" + ds._results_backend.close() + ds.conn.close() + + def test_raw_data_conn_is_set(self) -> None: + """When split is enabled, results go to a separate connection.""" + ds = new_data_set("test-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + # Once started, the results connection is the separate file, not the + # main database connection. + assert _uses_separate_file(ds) + assert ds._results_conn is not ds.conn + self._close_ds(ds) + + def test_sql_keyword_parameter_names_quoted_in_file(self) -> None: + """A parameter whose name is a SQL keyword (e.g. ``from``) must not + break creation of the per-dataset results table: startup must succeed + just as it does for the main-database backend, and the per-dataset file + must contain the quoted columns. + + (Inserting rows for such names is a separate, pre-existing QCoDeS + limitation shared by both backends, so it is not exercised here.) + """ + ds = new_data_set("test-split") + setpoint = ParamSpecBase("from", "numeric") + measured = ParamSpecBase("select", "numeric") + idps = InterDependencies_(dependencies={measured: (setpoint,)}) + ds.set_interdependencies(idps) + + # Previously raised sqlite3.OperationalError at CREATE TABLE time. + ds.mark_started() + + raw_conn = connect_to_raw_data_db(_raw_file(ds)) + try: + cursor = raw_conn.execute(f'PRAGMA table_info("{ds.table_name}")') + columns = {row[1] for row in cursor.fetchall()} + finally: + raw_conn.close() + assert {"from", "select"} <= columns + self._close_ds(ds) + + def test_background_writing_routes_to_separate_file(self) -> None: + """With split storage, background writes must land in the per-dataset + file (routed via the backend's results_db_path).""" + ds = new_data_set("test-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started(start_bg_writer=True) + + # The path the background writer targets is the separate file. + assert ds._results_backend.results_db_path is not None + assert Path(ds._results_backend.results_db_path) == _raw_file(ds) + + results = [{"x": float(i), "y": float(i**2)} for i in range(5)] + ds.add_results(results) + ds.mark_completed() # flushes the background writer + + assert ds.number_of_results == 5 + data = ds.get_parameter_data() + np.testing.assert_array_almost_equal( + data["y"]["y"], np.array([r["y"] for r in results]) + ) + # Nothing was written to the main DB (no results table there). + cursor = ds.conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name=?", + (ds.table_name,), + ) + assert cursor.fetchone() is None + self._close_ds(ds) + + def test_new_data_set_uses_separate_file_backend(self) -> None: + """When split is enabled, new datasets use the separate-file backend.""" + ds = new_data_set("test-split") + assert isinstance(ds, DataSet) + assert isinstance(ds._results_backend, SeparateSqliteFileResultsBackend) + self._close_ds(ds) + + def test_loaded_dataset_autodetects_backend(self) -> None: + """A split dataset is detected as such however it is loaded, including + via direct DataSet(run_id=...) construction.""" + ds, _ = self._make_dataset_with_data(n_rows=3) + run_id = ds.run_id + path_to_db = ds.path_to_db + self._close_ds(ds) + + # via the public load_by_id + loaded = load_by_id(run_id) + assert isinstance(loaded, DataSet) + assert isinstance(loaded._results_backend, SeparateSqliteFileResultsBackend) + self._close_ds(loaded) + + # via direct DataSet construction (auto-detection, no factory needed) + direct = DataSet(path_to_db=path_to_db, run_id=run_id) + assert isinstance(direct._results_backend, SeparateSqliteFileResultsBackend) + assert direct._results_conn is not direct.conn + self._close_ds(direct) + + def test_raw_data_file_created(self, tmp_path: Path) -> None: + """A per-dataset SQLite file should be created.""" + ds, _ = self._make_dataset_with_data() + raw_folder = tmp_path / "raw_data" + raw_file = raw_folder / f"{ds.guid}.db" + assert raw_file.is_file() + self._close_ds(ds) + + def test_raw_data_db_path_recorded(self) -> None: + """The path to the raw data file should be recorded in the runs table + (in a dedicated column) but kept out of the user-facing metadata.""" + ds, _ = self._make_dataset_with_data() + # Not exposed as user metadata ... + assert "raw_data_db_path" not in ds.metadata + # ... but recorded in the runs table and pointing at the real file. + assert _raw_file(ds).is_file() + assert _uses_separate_file(ds) + self._close_ds(ds) + + def test_data_is_in_raw_db_not_main(self) -> None: + """Data should be in the raw data file, and the main DB should not + contain a results table at all.""" + ds, _ = self._make_dataset_with_data(n_rows=5) + table_name = ds.table_name + main_conn = ds.conn + + # Main DB should NOT contain the results table - only metadata is kept + # there when raw data storage is enabled. + cursor = main_conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name=?", + (table_name,), + ) + assert cursor.fetchone() is None + + # The result_table_name identifier is still recorded in the runs table. + assert table_name == select_one_where( + main_conn, "runs", "result_table_name", "run_id", ds.run_id + ) + + # Raw data DB should have the actual data. The results connection is + # the separate file once the dataset has been started. + raw_conn = ds._results_conn + assert raw_conn is not ds.conn + cursor = raw_conn.execute(f'SELECT COUNT(*) FROM "{table_name}"') + raw_count = cursor.fetchone()[0] + assert raw_count == 5 + self._close_ds(ds) + + def test_number_of_results(self) -> None: + """number_of_results should read from the raw data file.""" + ds, _ = self._make_dataset_with_data(n_rows=7) + assert ds.number_of_results == 7 + self._close_ds(ds) + + def test_no_results_table_in_main_db_before_started(self) -> None: + """A pristine split dataset has no results table anywhere, and + counting results returns 0 without raising.""" + ds = new_data_set("test-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + + # Not started yet: the separate file is not created, so the results + # connection falls back to the main connection, and no table exists. + assert _uses_separate_file(ds) + assert ds._results_conn is ds.conn + cursor = ds.conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name=?", + (ds.table_name,), + ) + assert cursor.fetchone() is None + + # Counting must not raise even though the table does not exist. + assert ds.number_of_results == 0 + assert len(ds) == 0 + self._close_ds(ds) + + def test_subscribe_before_start_defers_and_fires(self) -> None: + """Subscribing before the dataset is started (before the raw data + table exists) should be deferred and then fire once data is added.""" + ds = new_data_set("test-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + + received: list[int] = [] + + def callback(results, length, state): + received.append(length) + + # Subscribe while still pristine - no results table exists yet. + sub_id = ds.subscribe(callback, min_wait=0, min_count=1) + assert sub_id in ds._pending_subscribers + assert ds.subscribers == {} + + ds.mark_started() + # Deferred subscriber has now been materialised. + assert ds._pending_subscribers == {} + assert sub_id in ds.subscribers + + ds.add_results([{"x": float(i), "y": float(i**2)} for i in range(5)]) + ds.mark_completed() + + assert len(received) > 0 + assert received[-1] == 5 + self._close_ds(ds) + + def test_get_parameter_data_from_raw_data(self) -> None: + """get_parameter_data should read from the raw data file.""" + ds, results = self._make_dataset_with_data(n_rows=5) + data = ds.get_parameter_data() + assert "y" in data + np.testing.assert_array_almost_equal( + data["y"]["x"], np.array([r["x"] for r in results]) + ) + np.testing.assert_array_almost_equal( + data["y"]["y"], np.array([r["y"] for r in results]) + ) + self._close_ds(ds) + + def test_cache_loads_from_raw_data(self) -> None: + """The cache should also read from the raw data file.""" + ds, results = self._make_dataset_with_data(n_rows=5) + cache = ds.cache + cache.load_data_from_db() + cache_data = cache.data() + assert "y" in cache_data + np.testing.assert_array_almost_equal( + cache_data["y"]["x"], np.array([r["x"] for r in results]) + ) + self._close_ds(ds) + + def test_load_by_id_with_split_data(self) -> None: + """Loading by ID should automatically use the raw data connection.""" + ds, results = self._make_dataset_with_data(n_rows=3) + run_id = ds.run_id + self._close_ds(ds) + + # Re-load from the database + loaded = load_by_id(run_id) + assert isinstance(loaded, DataSet) + assert _uses_separate_file(loaded) + assert loaded._results_conn is not loaded.conn + data = loaded.get_parameter_data() + np.testing.assert_array_almost_equal( + data["y"]["y"], np.array([r["y"] for r in results]) + ) + self._close_ds(loaded) + + def test_get_db_overview_counts_split_records(self) -> None: + """get_db_overview must count rows from the per-dataset raw-data file + for split runs (the results table is absent from the main DB).""" + ds, results = self._make_dataset_with_data(n_rows=7) + run_id = ds.run_id + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + overview = get_db_overview(path_to_db=db_path) + assert overview[run_id]["records"] == len(results) + + def test_multiple_datasets_split(self) -> None: + """Multiple datasets should each get their own raw data file.""" + ds1, _ = self._make_dataset_with_data(n_rows=3) + ds2, _ = self._make_dataset_with_data(n_rows=5) + + assert _uses_separate_file(ds1) + assert _uses_separate_file(ds2) + assert ds1._results_conn.path_to_dbfile != ds2._results_conn.path_to_dbfile + assert ds1.number_of_results == 3 + assert ds2.number_of_results == 5 + self._close_ds(ds1) + self._close_ds(ds2) + + def test_metadata_remains_in_main_db(self) -> None: + """Run metadata should be in the main DB, not the raw data DB.""" + ds, _ = self._make_dataset_with_data() + main_conn = ds.conn + # runs table should be populated in main DB + cursor = main_conn.execute( + "SELECT name, is_completed FROM runs WHERE run_id=?", (ds.run_id,) + ) + row = cursor.fetchone() + assert row is not None + assert row[0] == "test-split" + assert row[1] == 1 # completed + self._close_ds(ds) + + def test_missing_raw_data_file_raises(self, tmp_path: Path) -> None: + """Loading a dataset whose raw data file is missing should raise.""" + ds, _ = self._make_dataset_with_data(n_rows=3) + run_id = ds.run_id + raw_path = _raw_file(ds) + self._close_ds(ds) + + # Delete the raw data file + raw_path.unlink() + assert not raw_path.exists() + + # Attempting to load should raise FileNotFoundError + with pytest.raises(FileNotFoundError, match=r"Raw data file.*not found"): + load_by_id(run_id) + + +# --------------------------------------------------------------------------- +# Tests that the feature doesn't interfere when disabled +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("experiment") +class TestDataSetWithoutSplitRawData: + def test_uses_main_database_backend(self) -> None: + """When split is disabled, datasets use the main-database backend.""" + ds = new_data_set("test-no-split") + assert isinstance(ds._results_backend, MainDatabaseResultsBackend) + # Results go to the main database connection. + assert ds._results_conn is ds.conn + ds.conn.close() + + def test_results_conn_is_main_conn(self) -> None: + """When split is disabled, the results connection is the main one.""" + ds = new_data_set("test-no-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + assert ds._results_conn is ds.conn + ds.conn.close() + + def test_data_in_main_db(self) -> None: + """When split is disabled, data should be in the main DB as usual.""" + ds = new_data_set("test-no-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + + cursor = ds.conn.execute(f'SELECT COUNT(*) FROM "{ds.table_name}"') + assert cursor.fetchone()[0] == 1 + assert ds.number_of_results == 1 + ds.conn.close() + + +# --------------------------------------------------------------------------- +# Tests for update_raw_data_paths helper +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_raw_data_experiment") +class TestUpdateRawDataPaths: + @staticmethod + def _close_ds(ds: DataSet) -> None: + ds._results_backend.close() + ds.conn.close() + + def test_update_after_move(self, tmp_path: Path) -> None: + """Paths should be updated after raw data files are moved.""" + ds = new_data_set("test-move") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + + raw_path = _raw_file(ds) + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + # Move the raw data file to a new folder + new_folder = tmp_path / "moved_raw" + new_folder.mkdir() + new_file = new_folder / raw_path.name + + shutil.move(str(raw_path), str(new_file)) + + # Run the update + updated = update_raw_data_paths(db_path, new_folder) + assert len(updated) == 1 + run_id, old_path, new_path = updated[0] + assert old_path == str(raw_path) + assert new_path == str(new_file) + + # Verify the dataset can now be loaded + loaded = load_by_id(run_id) + assert isinstance(loaded, DataSet) + assert loaded.number_of_results == 1 + data = loaded.get_parameter_data() + assert data["y"]["x"][0] == 1.0 + self._close_ds(loaded) + + def test_update_skips_missing_files(self, tmp_path: Path) -> None: + """Runs whose files are not in the new folder should be skipped.""" + ds = new_data_set("test-skip") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + # Point to an empty folder — file not there + empty_folder = tmp_path / "empty" + empty_folder.mkdir() + + updated = update_raw_data_paths(db_path, empty_folder) + assert len(updated) == 0 + + def test_update_nonexistent_db_raises(self, tmp_path: Path) -> None: + """Should raise if the main DB file doesn't exist.""" + with pytest.raises(FileNotFoundError, match="Database file not found"): + update_raw_data_paths(tmp_path / "nonexistent.db", tmp_path) + + def test_update_nonexistent_folder_raises(self, tmp_path: Path) -> None: + """Should raise if the new folder doesn't exist.""" + db_path = tmp_path / "test.db" + db_path.touch() + with pytest.raises(FileNotFoundError, match="New raw data folder not found"): + update_raw_data_paths(db_path, tmp_path / "no_such_folder") + + +# --------------------------------------------------------------------------- +# purge_orphaned_datasets and cleanup_datasets tests +# --------------------------------------------------------------------------- + + +class TestPurgeOrphanedDatasets: + """Tests for purge_orphaned_datasets.""" + + @staticmethod + def _close_ds(ds: DataSet) -> None: + ds._results_backend.close() + ds.conn.close() + + @pytest.mark.usefixtures("_raw_data_db") + def test_dry_run_reports_orphans(self, tmp_path: Path) -> None: + """Dry run should report orphaned datasets without removing them.""" + + new_experiment("test-exp", sample_name="test-sample") + ds = new_data_set("test-purge") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + + raw_path = _raw_file(ds) + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + # Delete the raw data file to create an orphan + raw_path.unlink() + + result = purge_orphaned_datasets(db_path, dry_run=True) + assert result.total_datasets_with_raw_data == 1 + assert len(result.orphaned_datasets) == 1 + assert len(result.removed_datasets) == 0 + assert result.dry_run is True + + # Verify the dataset is still in the DB (direct query since load_by_id + # will raise FileNotFoundError due to missing raw data file) + + conn = connect(db_path) + cursor = conn.execute("SELECT COUNT(*) FROM runs WHERE run_id = 1") + assert cursor.fetchone()[0] == 1 + conn.close() + + @pytest.mark.usefixtures("_raw_data_db") + def test_removes_orphans_when_not_dry_run(self, tmp_path: Path) -> None: + """With dry_run=False, orphaned datasets should be removed.""" + + new_experiment("test-exp", sample_name="test-sample") + ds = new_data_set("test-purge-remove") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + + raw_path = _raw_file(ds) + db_path = ds.path_to_db + assert db_path is not None + run_id = ds.run_id + self._close_ds(ds) + + # Delete the raw data file + raw_path.unlink() + + result = purge_orphaned_datasets(db_path, dry_run=False) + assert len(result.orphaned_datasets) == 1 + assert len(result.removed_datasets) == 1 + assert result.removed_datasets[0].run_id == run_id + assert result.dry_run is False + + # Verify the dataset no longer exists in the DB + + conn = connect(db_path) + cursor = conn.execute("SELECT COUNT(*) FROM runs WHERE run_id = ?", (run_id,)) + assert cursor.fetchone()[0] == 0 + conn.close() + + @pytest.mark.usefixtures("_raw_data_db") + def test_skips_datasets_with_existing_files(self, tmp_path: Path) -> None: + """Datasets whose raw data files exist should not be purged.""" + + new_experiment("test-exp", sample_name="test-sample") + ds = new_data_set("test-no-purge") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + # Don't delete the file — should not be orphaned + result = purge_orphaned_datasets(db_path, dry_run=False) + assert len(result.orphaned_datasets) == 0 + assert len(result.removed_datasets) == 0 + + def test_nonexistent_db_raises(self, tmp_path: Path) -> None: + """Should raise if the DB file doesn't exist.""" + + with pytest.raises(FileNotFoundError, match="Database file not found"): + purge_orphaned_datasets(tmp_path / "nonexistent.db") + + +class TestCleanupDatasets: + """Tests for cleanup_datasets.""" + + @staticmethod + def _close_ds(ds: DataSet) -> None: + ds._results_backend.close() + ds.conn.close() + + @pytest.mark.usefixtures("_raw_data_db") + def test_cleanup_by_sample_name(self, tmp_path: Path) -> None: + """Should remove datasets matching exact sample name.""" + + new_experiment("exp1", sample_name="sample-A") + ds1 = new_data_set("ds1") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds1.set_interdependencies(idps) + ds1.mark_started() + ds1.add_results([{"x": 1.0, "y": 2.0}]) + ds1.mark_completed() + raw_path1 = _raw_file(ds1) + db_path = ds1.path_to_db + assert db_path is not None + self._close_ds(ds1) + + # Create another dataset with a different sample + new_experiment("exp2", sample_name="sample-B") + ds2 = new_data_set("ds2") + ds2.set_interdependencies(idps) + ds2.mark_started() + ds2.add_results([{"x": 3.0, "y": 4.0}]) + ds2.mark_completed() + raw_path2 = _raw_file(ds2) + self._close_ds(ds2) + + # Dry run first + result = cleanup_datasets(db_path, sample_name="sample-A", dry_run=True) + assert len(result.matching_datasets) == 1 + assert len(result.removed_datasets) == 0 + assert raw_path1.is_file() + + # Actual removal + result = cleanup_datasets(db_path, sample_name="sample-A", dry_run=False) + assert len(result.matching_datasets) == 1 + assert len(result.removed_datasets) == 1 + assert not raw_path1.is_file() + assert raw_path2.is_file() # other sample untouched + + @pytest.mark.usefixtures("_raw_data_db") + def test_cleanup_skips_in_progress_runs(self, tmp_path: Path) -> None: + """An in-progress (started, not completed) run must never be removed, + even if it matches the cleanup criteria - its file may still be + actively written.""" + new_experiment("exp1", sample_name="test-sample") + ds = new_data_set("ds-active") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + ds.set_interdependencies(InterDependencies_(dependencies={y: (x,)})) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + # deliberately NOT marked completed + raw_path = _raw_file(ds) + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + result = cleanup_datasets(db_path, sample_name="test-sample", dry_run=False) + assert len(result.matching_datasets) == 0 + assert len(result.removed_datasets) == 0 + assert raw_path.is_file() + + @pytest.mark.usefixtures("_raw_data_db") + def test_cleanup_by_size(self, tmp_path: Path) -> None: + """Should remove datasets larger than the specified threshold.""" + + new_experiment("exp1", sample_name="test-sample") + ds = new_data_set("ds-large") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + # Add enough data so the file has some size + for i in range(100): + ds.add_results([{"x": float(i), "y": float(i * 2)}]) + ds.mark_completed() + raw_path = _raw_file(ds) + db_path = ds.path_to_db + assert db_path is not None + self._close_ds(ds) + + file_size_mb = raw_path.stat().st_size / (1024 * 1024) + + # Set threshold below actual size — should match + result = cleanup_datasets( + db_path, larger_than_mb=file_size_mb * 0.5, dry_run=True + ) + assert len(result.matching_datasets) == 1 + + # Set threshold above actual size — should not match + result = cleanup_datasets( + db_path, larger_than_mb=file_size_mb * 2, dry_run=True + ) + assert len(result.matching_datasets) == 0 + + @pytest.mark.usefixtures("_raw_data_db") + def test_cleanup_by_age(self, tmp_path: Path) -> None: + """Should remove datasets older than the specified number of days.""" + + new_experiment("exp1", sample_name="test-sample") + ds = new_data_set("ds-old") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + ds.add_results([{"x": 1.0, "y": 2.0}]) + ds.mark_completed() + db_path = ds.path_to_db + assert db_path is not None + run_id = ds.run_id + self._close_ds(ds) + + # Set the completed_timestamp to 10 days ago + old_ts = time.time() - (10 * 86400) + conn = connect(db_path) + conn.execute( + "UPDATE runs SET completed_timestamp = ? WHERE run_id = ?", + (old_ts, run_id), + ) + conn.commit() + conn.close() + + # older_than_days=5 should match (ds is 10 days old) + result = cleanup_datasets(db_path, older_than_days=5, dry_run=True) + assert len(result.matching_datasets) == 1 + + # older_than_days=15 should NOT match (ds is only 10 days old) + result = cleanup_datasets(db_path, older_than_days=15, dry_run=True) + assert len(result.matching_datasets) == 0 + + def test_no_criteria_raises(self, tmp_path: Path) -> None: + """Should raise if no criteria are specified.""" + + db_path = tmp_path / "test.db" + db_path.touch() + with pytest.raises(ValueError, match="At least one cleanup criterion"): + cleanup_datasets(db_path) + + def test_negative_older_than_days_raises(self, tmp_path: Path) -> None: + """A negative age threshold would match (and delete) everything.""" + db_path = tmp_path / "test.db" + db_path.touch() + with pytest.raises(ValueError, match="older_than_days must be non-negative"): + cleanup_datasets(db_path, older_than_days=-1, dry_run=False) + + def test_negative_larger_than_mb_raises(self, tmp_path: Path) -> None: + """A negative size threshold would match (and delete) everything.""" + db_path = tmp_path / "test.db" + db_path.touch() + with pytest.raises(ValueError, match="larger_than_mb must be non-negative"): + cleanup_datasets(db_path, larger_than_mb=-1, dry_run=False) + + def test_nonexistent_db_raises(self, tmp_path: Path) -> None: + """Should raise if the DB file doesn't exist.""" + + with pytest.raises(FileNotFoundError, match="Database file not found"): + cleanup_datasets(tmp_path / "nonexistent.db", older_than_days=1) + + +# --------------------------------------------------------------------------- +# Integration tests - extract/export with split raw data +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_raw_data_experiment") +class TestExtractExportWithSplitRawData: + """Verify that extract_runs_into_db and export_datasets_and_create_metadata_db + work correctly when raw data is stored in separate per-dataset SQLite files.""" + + @staticmethod + def _make_dataset_with_data( + n_rows: int = 10, + ) -> tuple[DataSet, list[dict[str, float]]]: + ds = new_data_set("test-split") + x = ParamSpecBase("x", "numeric") + y = ParamSpecBase("y", "numeric") + idps = InterDependencies_(dependencies={y: (x,)}) + ds.set_interdependencies(idps) + ds.mark_started() + results = [{"x": float(i), "y": float(i**2)} for i in range(n_rows)] + ds.add_results(results) + ds.mark_completed() + return ds, results + + @staticmethod + def _close_ds(ds: DataSet) -> None: + ds._results_backend.close() + ds.conn.close() + + def test_extract_runs_into_db_with_split_data(self, tmp_path: Path) -> None: + """extract_runs_into_db should copy data from external raw data file.""" + ds, results = self._make_dataset_with_data(n_rows=5) + run_id = ds.run_id + source_db = ds.path_to_db + assert source_db is not None + self._close_ds(ds) + + target_db = str(tmp_path / "target.db") + extract_runs_into_db(source_db, target_db, run_id) + + # Verify data was copied to target DB + target_conn = connect(target_db) + target_ds = load_by_id(1, conn=target_conn) + data = target_ds.get_parameter_data() + np.testing.assert_array_almost_equal( + data["y"]["y"], np.array([r["y"] for r in results]) + ) + assert target_ds.number_of_results == 5 + # The extracted dataset owns its data in the target DB and must not + # keep referring to the source's per-dataset raw data file. + assert get_raw_data_db_path_for_run(target_conn, target_ds.run_id) is None + assert "raw_data_db_path" not in target_ds.metadata + target_conn.close() + + def test_export_and_create_metadata_db_with_split_data( + self, tmp_path: Path + ) -> None: + """export_datasets_and_create_metadata_db should work with split data.""" + ds, _results = self._make_dataset_with_data(n_rows=5) + source_db = ds.path_to_db + assert source_db is not None + self._close_ds(ds) + + target_db = tmp_path / "metadata_only.db" + export_path = tmp_path / "netcdf_exports" + + statuses = export_datasets_and_create_metadata_db( + source_db, target_db, export_path=export_path + ) + + assert len(statuses) == 1 + # Should either be exported to NetCDF or copied as-is + assert next(iter(statuses.values())) in ("exported", "copied_as_is")