diff --git a/.gitignore b/.gitignore index 3425b74a3b..03bef45d23 100644 --- a/.gitignore +++ b/.gitignore @@ -70,6 +70,7 @@ magic.odm.yaml # This directory contains the database from fluorophores.org # To fill it, run ./util/getfluodb.py /install/linux/usr/share/odemis/fluodb/ +/install/linux/usr/share/odemis/datacollector.key # External drivers /src/odemis/driver/hitachi.py diff --git a/debian/control b/debian/control index 164bb886f9..ff920b5ab1 100644 --- a/debian/control +++ b/debian/control @@ -75,6 +75,8 @@ Depends: ${shlibs:Depends}, python3-libusb1, # Needed for acq.fastem python3-shapely, +# Needed for communicating with AWS S3 + python3-boto3, Suggests: imagej Description: Open Delmic Microscope Software Odemis is the acquisition software for the Delmic microscopes. In particular, diff --git a/debian/odemis.install b/debian/odemis.install index 13eb336a96..6384ba6270 100644 --- a/debian/odemis.install +++ b/debian/odemis.install @@ -1,5 +1,6 @@ install/linux/etc/logrotate.d/odemisd etc/logrotate.d/ install/linux/etc/sudoers.d/odemis etc/sudoers.d/ +install/linux/usr/share/odemis/datacollector.key usr/share/odemis/ install/linux/usr/share/bash-completion/completions/odemis-cli usr/share/bash-completion/completions/ install/linux/usr/share/bash-completion/completions/odemis-convert usr/share/bash-completion/completions/ install/linux/usr/share/fish/vendor_completions.d/odemis.fish usr/share/fish/vendor_completions.d/ diff --git a/debian/rules b/debian/rules index ee4213b60d..b88f03f9d3 100755 --- a/debian/rules +++ b/debian/rules @@ -10,6 +10,9 @@ export PYBUILD_DISABLE=test dh $@ --with python3 --with sphinxdoc --buildsystem=pybuild override_dh_auto_install: + @if [ ! -f install/linux/usr/share/odemis/datacollector.key ]; then \ + echo "WARNING: Missing install/linux/usr/share/odemis/datacollector.key"; \ + fi python3 setup.py install --root debian/odemis --install-layout=deb override_dh_installdocs: diff --git a/doc/develop/data-framework-setup-guide.rst b/doc/develop/data-framework-setup-guide.rst new file mode 100644 index 0000000000..6ec76a4271 --- /dev/null +++ b/doc/develop/data-framework-setup-guide.rst @@ -0,0 +1,362 @@ +Odemis Data Collection Framework — Setup Guide +=============================================== + +**Audience:** Software and support engineers setting up data collection on an +Odemis installation or on a developer workstation. + +.. note:: + All commands are Bash. Enter them line by line unless stated otherwise. + +This guide covers: + +- Prerequisites +- Installing the boto3 package +- Creating and installing the credentials key file +- Setting up the local queue directory (``dc_queue``) +- Verifying the configuration file +- Test-bucket vs production-bucket mode +- Running the unit tests +- Quick smoke-test with ``odemis-dc-fetch`` +- Troubleshooting + + +Prerequisites +------------- + +- Ubuntu 22.04 LTS or later +- Odemis installed from the Debian package or a source checkout +- AWS S3 credentials (``access_key`` + ``secret_key``) for the target bucket + (typically upload-only IAM keys). + Obtain these from a Delmic software engineer; they are **not** stored in this + repository. +- Write access to ``/usr/share/odemis/`` (for the key file) + + +Installing the boto3 Package +----------------------------- + +``boto3`` is the AWS SDK for Python used to upload data to S3. + +When installed via the Odemis Debian package it is listed as a dependency and +will be pulled in automatically: + +.. code-block:: bash + + sudo apt update + sudo apt install python3-boto3 + +To verify the installation: + +.. code-block:: bash + + python3 -c "import boto3; print(boto3.__version__)" + +If the command prints a version string (e.g. ``1.28.0``) the package is ready. +If it raises ``ImportError``, re-run the ``apt install`` command above. + + +Installing the Credentials Key File +------------------------------------- + +The framework reads S3 credentials from a single JSON file at: /usr/share/odemis/datacollector.key + +The file must contain exactly the following two keys: + +.. code-block:: json + + { "access_key": "", "secret_key": "" } + +Obtain the actual key values from a Delmic software engineer. These keys are +typically scoped to upload-only access (no read/retrieve), so internal sharing +over standard channels is lower risk than full-access AWS credentials. + +Do **not** commit these keys to git repositories (especially public remotes). +Public repository scrapers continuously harvest AWS key patterns, and committing +keys can also trigger GitHub secret-scanning/security warnings. + + +Production Setup +~~~~~~~~~~~~~~~~ + +Contact a software engineer who has access to the AWS IAM console to get the +key pair, then run the following commands line by line: + +.. code-block:: bash + + # Create the file — replace placeholder values with the real keys + sudo tee /usr/share/odemis/datacollector.key > /dev/null << 'EOF' + { "access_key": "", "secret_key": "" } + EOF + + +Setting up the Local Queue Directory (dc_queue) +------------------------------------------------- + +.. note:: + This section applies to the **test / developer setup** only. Skip it for a + production installation. + +The framework stages serialised ZIP archives in a local directory before +uploading them. The default path is: + +.. code-block:: text + + ~/.local/share/odemis/dc_queue + +This directory is created automatically by the framework at runtime if it does +not exist, as long as the parent ``~/.local/share/odemis/`` is writable by the +process. + +For a standard Odemis installation ``~/.local/share/odemis/`` is already +created by the package post-install script. If that is not the case, create it +manually: + +.. code-block:: bash + + sudo mkdir -p ~/.local/share/odemis/dc_queue + sudo chown $USER:$USER ~/.local/share/odemis/dc_queue + sudo chmod 750 ~/.local/share/odemis/dc_queue + + +Queue Disk Limit +~~~~~~~~~~~~~~~~ + +The framework automatically enforces a soft limit of 10 % of the partition's +total disk space on ``~/.local/share/odemis/dc_queue``. When the limit is +exceeded, the oldest ZIP files are deleted with a ``WARNING`` log entry. + +- **Production:** ``~/.local/share/odemis/`` is normally on the main OS + partition. No special configuration is required. +- **Test / developer workstation:** the queue directory may be redirected to a + temporary location by instantiating ``_BackgroundWorker`` with a custom + ``queue_dir`` argument (used in unit tests). For manual testing with the real + Odemis GUI, the default path is always used. + + +Inspecting Queued Files +~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + ls -lh ~/.local/share/odemis/dc_queue/ + +Files follow the naming convention:: + + --.zip + +Any file ending in ``.tmp`` is an incomplete write left over from a crash; the +framework removes such files on startup. You can delete them manually if Odemis +is not running: + +.. code-block:: bash + + rm -f ~/.local/share/odemis/dc_queue/*.tmp + + +Verifying the Configuration File +---------------------------------- + +The per-user consent state is stored in: + +.. code-block:: text + + ~/.config/odemis/datacollector.config + +This file is created automatically the first time the user interacts with the +consent dialog from the menu. Its permissions are always set to ``0600`` +(owner read/write only) by the framework. + +A typical file after opt-in looks like: + +.. code-block:: ini + + [general] + # Data sharing consent (true / false). + consent = true + + +.. _test-bucket-mode: + +Test-Bucket vs Production-Bucket Mode +--------------------------------------- + +Production (Default) +~~~~~~~~~~~~~~~~~~~~~ + +- **Bucket:** ``delmic-odemis-collect`` +- **Region:** ``eu-west-1`` + +This is the default when the framework starts normally. No environment variable +needs to be set. Use the production IAM credentials in the key file +(see `Installing the Credentials Key File`_). + + +Test / Developer Mode +~~~~~~~~~~~~~~~~~~~~~~ + +- **Bucket:** ``delmic-odemis-collect-test`` +- **Region:** ``eu-west-1`` + +Set the following environment variable before starting Odemis (or any script +that calls ``record()``): + +.. code-block:: bash + + export TEST_DATACOLLECTION=1 + +The framework logs the following ``INFO`` message when test mode is active:: + + DataCollector: TEST_DATACOLLECTION=1 — using test bucket 'delmic-odemis-collect-test' + +.. important:: + The test bucket credentials (``access_key`` / ``secret_key`` in the key + file) must be scoped to the **test** bucket by the Delmic AWS admin. If you + install production credentials and set ``TEST_DATACOLLECTION=1``, uploads + will fail with an ``AccessDenied`` error because the IAM policy only permits + writes to the production prefix. + + +Managing Consent +---------------- + +The data collection consent dialog is accessed via the Odemis menu. Users have +three options: + +- **Opt In:** Enable data collection permanently. +- **Opt Out:** Disable data collection permanently. +- **Consent for One Day:** Enable data collection temporarily for one day only. + After the specified day, consent automatically expires to disabled. + +When temporary consent is active and less than one day remains, the collection +sampling rate increases to 100% to prioritize data collection. Otherwise, a +default 10% sampling rate applies. + +When consent is active, collected data is staged in the queue directory and +uploaded to S3 in the background. When consent is disabled, no data is collected +or uploaded. + + +Running the Unit Tests +----------------------- + +Unit tests run without any hardware and without a real S3 connection. All +upload calls are mocked. + +.. code-block:: bash + + # From the repository root + env TEST_NOHW=1 python3 src/odemis/util/test/datacollector_test.py + +To run a specific test class or method: + +.. code-block:: bash + + env TEST_NOHW=1 python3 src/odemis/util/test/datacollector_test.py \ + DataCollectorTest.test_record_returns_fast + + env TEST_NOHW=1 python3 src/odemis/util/test/datacollector_test.py \ + TestSerialize.test_metadata_json_envelope_fields + + +Real S3 Integration Tests +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The class ``TestRealS3Integration`` uploads to the test bucket and cleans up +after itself. It requires: + +- ``boto3`` installed +- ``/usr/share/odemis/datacollector.key`` present with test bucket credentials +- ``TEST_DATACOLLECTION=1`` is **not** needed here — the test class hard-codes + the test bucket directly + +.. code-block:: bash + + env TEST_NOHW=1 python3 src/odemis/util/test/datacollector_test.py \ + TestRealS3Integration + +If the key file is absent, this class is skipped automatically. + + +Quick Smoke-Test with odemis-dc-fetch +--------------------------------------- + +After a successful upload (production or test), use the retrieval script to +confirm objects landed in S3. The retrieval script can only be used with an AWS +profile that has data-analyst read access; the upload key in +``/usr/share/odemis/datacollector.key`` is typically write-only and usually +does not grant retrieval rights. + +.. code-block:: bash + + odemis-dc-fetch \ + --bucket delmic-odemis-collect \ + --region eu-west-1 \ + --output ./dc_samples_test + +Filter by event name and date: + +.. code-block:: bash + + odemis-dc-fetch \ + --bucket delmic-odemis-collect-test \ + --region eu-west-1 \ + --event feature_collected \ + --since 2026-04-01 \ + --output ./dc_samples_test + +The script prints a one-line summary:: + + listed=N matched=N downloaded=N skipped_existing=N failed=0 + +If ``failed`` is non-zero, check the log output for ``AccessDenied`` or network +errors and verify the key file credentials. + + +Troubleshooting +---------------- + +**Problem:** ``LookupError: S3 credentials key file not found at /usr/share/odemis/datacollector.key`` + +**Solution:** Create the key file as described in +`Installing the Credentials Key File`_. + +---- + +**Problem:** ``botocore.exceptions.ClientError: AccessDenied`` + +**Cause:** The IAM key in the key file does not have permission to write to the +bucket being targeted. + +**Solution:** Ensure the key file contains credentials matching the target +bucket (production key → production bucket, test key → test bucket). Check +whether ``TEST_DATACOLLECTION=1`` is set unexpectedly. + +---- + +**Problem:** Uploads never happen; queue fills up. + +**Cause:** Network is unavailable or credentials are wrong. + +**Solution:** Check the Odemis log for ``DataCollector upload failed`` entries. +The framework retries with exponential back-off (30 s → 60 s → ... up to 1 h). +Pending ZIPs remain in ``~/.local/share/odemis/dc_queue/`` and are flushed +oldest-first once connectivity is restored. + +---- + +**Problem:** ``Queue limit exceeded: removed oldest sample`` appears in the log. + +**Cause:** The queue directory has grown beyond 10 % of the partition. + +**Solution:** Check disk space with ``df -h ~/.local/share/odemis/dc_queue``. +Investigate why uploads are not succeeding (credentials, network). + +---- + +**Problem:** ``TEST_DATACOLLECTION=1`` is set but uploads still go to +production. + +**Solution:** The environment variable is read at the +moment ``get_upload_backend()`` is called, which is lazy (first upload attempt). +Verify the variable is exported in the same shell/environment that runs the +Odemis process. diff --git a/doc/develop/index.rst b/doc/develop/index.rst index f32e9ae8ea..7832c8914d 100644 --- a/doc/develop/index.rst +++ b/doc/develop/index.rst @@ -25,6 +25,7 @@ This document covers the main aspects of development with Odemis. gui plugins utils + data-framework-setup-guide microscope documentation .. testing diff --git a/src/odemis/gui/cont/menu.py b/src/odemis/gui/cont/menu.py index 8d8da61956..b50d5ef31e 100644 --- a/src/odemis/gui/cont/menu.py +++ b/src/odemis/gui/cont/menu.py @@ -27,7 +27,9 @@ import odemis.gui.conf from odemis.gui.model import CHAMBER_VACUUM, CHAMBER_UNKNOWN from odemis.gui.model.dye import DyeDatabase +from odemis.gui.model.main_gui_data import MainGUIData from odemis.gui.util import call_in_wx_main +from odemis.util.datacollector import DataCollector from odemis.util import driver import os import subprocess @@ -42,14 +44,17 @@ class MenuController(object): tab controller. """ - def __init__(self, main_data, main_frame): - """ Binds the menu actions. + def __init__(self, main_data: MainGUIData, main_frame: wx.Frame, data_collector: DataCollector): + """ + Binds the menu actions. - main_data (MainGUIData): the representation of the microscope GUI - main_frame: (wx.Frame): the main frame of the GUI + :param main_data: The representation of the microscope GUI. + :param main_frame: The main frame of the GUI. + :param data_collector: The data collector, used for the data sharing consent menu item. """ self._main_data = main_data self._main_frame = main_frame + self._data_collector = data_collector # /File # /File/Open... @@ -153,6 +158,9 @@ def __init__(self, main_data, main_frame): # /Help/About main_frame.Bind(wx.EVT_MENU, self._on_about, id=main_frame.menu_item_about.GetId()) + self._consent_menu_item = self._append_data_sharing_menu_item(main_frame) + if self._consent_menu_item is not None: + main_frame.Bind(wx.EVT_MENU, self._on_toggle_data_sharing, id=self._consent_menu_item.GetId()) # add a toggle for correlation tab in viewer mode if main_data.is_viewer: @@ -163,6 +171,56 @@ def __init__(self, main_data, main_frame): menu.Remove(main_frame.menu_item_show_correlation) main_frame.menu_item_show_correlation.Destroy() + def _append_data_sharing_menu_item(self, main_frame: wx.Frame) -> wx.MenuItem: + """ + Append and initialize Help menu checkbox for data sharing consent. + :param main_frame: The main application frame. + :return: The created menu item, or None if the Help menu is not available. + """ + help_menu = main_frame.menu_item_about.GetMenu() + if help_menu is None: + return None + help_menu.AppendSeparator() + item = help_menu.AppendCheckItem(wx.ID_ANY, "Share data with Delmic") + item.Check(self._data_collector.get_consent() is True) + return item + + @call_in_wx_main + def _on_toggle_data_sharing(self, evt): + """Show consent dialog when the data-sharing menu item is clicked.""" + try: + message = ("Allow Odemis to send selected acquisition data to Delmic, " + "including example images and interaction logs (e.g., alignment steps, point-of-interest placements). " + "A random sample (eg, 10%) of workflow data may be included. This data is used solely to improve " + "Odemis algorithms and user experience. Data is stored securely on EU servers, accessible " + "only to designated Delmic analysts, and never shared with third parties.\n\n" + "By opting in, you confirm you have the authority to make this decision for your institution" + " and that sharing this data is permitted under your institution's data policies.\n\n" + "\"Enable for one day\" enables sharing for 24 hours with 100% collection, " + "then automatically disables it.") + + dlg = wx.MessageDialog( + self._main_frame, + message=message, + caption="Share data with Delmic", + style=wx.YES_NO | wx.CANCEL | wx.ICON_QUESTION + ) + dlg.SetYesNoCancelLabels(yes="Opt in", no="Opt out", cancel="Enable for one day") + response = dlg.ShowModal() + dlg.Destroy() + + if response == wx.ID_YES: + self._data_collector.set_consent(True) + elif response == wx.ID_NO: + self._data_collector.set_consent(False) + elif response == wx.ID_CANCEL: + self._data_collector.set_temporary_consent(days=1) + # Sync the Help menu checkbox to reflect the persisted choice. + if self._consent_menu_item is not None: + self._consent_menu_item.Check(self._data_collector.get_consent() is True) + except Exception: + logging.exception("Failed to run data-collection consent prompt.") + def _on_update(self, evt): import odemis.gui.util.updater as updater u = updater.WindowsUpdater() diff --git a/src/odemis/gui/main.py b/src/odemis/gui/main.py index cb5bf09418..5415c9649a 100755 --- a/src/odemis/gui/main.py +++ b/src/odemis/gui/main.py @@ -32,6 +32,7 @@ from odemis.gui.cont.temperature import TemperatureController from odemis.gui.util import call_in_wx_main from odemis.gui.xmlh import odemis_get_resources +from odemis.util.datacollector import DataCollector import sys import threading import traceback @@ -73,6 +74,7 @@ def __init__(self, standalone=False, file_name=None): self._snapshot_controller = None self._temperature_controller = None self._menu_controller = None + self._data_collector = DataCollector() self.plugins = [] # List of instances of plugin.Plugins # User input devices @@ -344,7 +346,7 @@ def toggle_log_panel(_): self.main_data.level.subscribe(self.on_level_va, init=True) log.create_gui_logger(self.main_frame.txt_log, self.main_data.debug, self.main_data.level) - self._menu_controller = MenuController(self.main_data, self.main_frame) + self._menu_controller = MenuController(self.main_data, self.main_frame, self._data_collector) # Menu events self.main_frame.Bind(wx.EVT_MENU, self.on_close_window, id=self.main_frame.menu_item_quit.GetId()) @@ -375,6 +377,7 @@ def toggle_log_panel(_): # Due to a bug in wxPython, sometimes the .Maximize() at the beginning of the function # has no effect. So we call it after the Show() to be sure it works. wx.CallAfter(self.main_frame.Maximize) + except Exception: self.excepthook(*sys.exc_info()) # Re-raise the exception, so the program will exit. If this is not diff --git a/src/odemis/util/datacollector.py b/src/odemis/util/datacollector.py new file mode 100644 index 0000000000..df00016df0 --- /dev/null +++ b/src/odemis/util/datacollector.py @@ -0,0 +1,812 @@ +# -*- coding: utf-8 -*- +""" +Created on 11 March 2026 + +@author: Karishma Kumar + +Copyright © 2026 Karishma Kumar, Delmic + +This file is part of Odemis. + +Odemis is free software: you can redistribute it and/or modify it under the +terms of the GNU General Public License version 2 as published by the Free +Software Foundation. + +Odemis is distributed in the hope that it will be useful, but WITHOUT ANY +WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR +A PARTICULAR PURPOSE. See the GNU General Public License for more details. + +You should have received a copy of the GNU General Public License along with +Odemis. If not, see http://www.gnu.org/licenses/. + +Odemis Annotated Data Collection Framework. + +Provides a thread-safe, non-blocking DataCollector.record() call that any +Odemis module can invoke to capture a labelled data sample. Serialisation +happens asynchronously in a background daemon thread; the caller returns +immediately. +""" + +import configparser +import json +import logging +import os +import queue +import random +import re +import shutil +import socket +import tempfile +import threading +import time +import uuid +import zipfile +from dataclasses import dataclass, field +from datetime import date, datetime, timedelta, timezone +from pathlib import Path +from typing import Optional + +try: + import boto3 +except ImportError: + logging.error("boto3 is required for S3 upload functionality; install with 'sudo apt install python3-boto3'") + raise +import numpy + +import odemis +from odemis import model +from odemis.dataio import hdf5, tiff + + +# S3 bucket name — shared production bucket, created once by the dev team. +S3_BUCKET = "delmic-odemis-collect" + +# S3 bucket used for automated tests (not the production bucket). +S3_TEST_BUCKET = "delmic-odemis-collect-test" + +# S3 endpoint URL — None means let boto3 resolve the regional endpoint automatically. +# Set explicitly only for custom S3-compatible storage. +S3_ENDPOINT_URL = None +S3_REGION = "eu-west-1" + +# Path to the S3 credentials key file (JSON with access_key / secret_key). +_CREDENTIALS_PATH = "/usr/share/odemis/datacollector.key" + +# Default paths +_CONF_DIR = os.path.join(os.path.expanduser("~"), ".config", "odemis") +_DEFAULT_QUEUE_DIR = Path("~/.local/share/odemis/dc_queue") + +_VALID_IMAGE_FORMATS = ("TIFF", "HDF5") +_INITIAL_RETRY_DELAY_SECONDS = 30.0 +_MAX_RETRY_DELAY_SECONDS = 3600.0 +CONSENT_DATE_KEY = "consent_date" + +# Collection probability: fraction of record() calls that are actually uploaded. +# 100% when consent expires within 1 day (1-day trial), 10% otherwise. +_DEFAULT_COLLECTION_PROBABILITY = 0.10 +_FULL_COLLECTION_PROBABILITY = 1.0 + + +def _sanitize_filename(name: str) -> str: + """Return a filesystem-safe version of name. + + Strips path components (prevents traversal) and replaces any character + that is not alphanumeric, -, _ or . with an underscore. + + :param name: Raw name to sanitize. + :returns: Safe filename string; never empty (falls back to "_"). + """ + # Strip path components to prevent directory traversal. + name = os.path.basename(name.replace("\\", "/")) + # Replace characters unsafe in filenames. + name = re.sub(r"[^\w\-.]", "_", name) + return name or "_" + +def _search_credentials() -> dict: + """ + Load S3 credentials from the standard key-file location. + The key file is a JSON file containing access_key and secret_key. + :returns: Dict with access_key and secret_key. + :raises LookupError: If the key file is not found at the expected location. + """ + if not os.path.isfile(_CREDENTIALS_PATH): + raise LookupError( + f"S3 credentials key file not found at {_CREDENTIALS_PATH}" + ) + with open(_CREDENTIALS_PATH, "r") as fh: + data = json.load(fh) + return { + "access_key": data["access_key"], + "secret_key": data["secret_key"], + } + + +class DataCollectorConfig: + """Persistent configuration for the data-collection framework. + + Backed by a configparser INI file at + ~/.config/odemis/datacollector.config. + + Sections + -------- + [general] + consent — true / false / none (not yet decided). + consent_date — Date (ISO UTC) until which consent is active. + Consent auto-expires to False after this date. + Commented out when not applicable. + + The file is written in a human-readable format with inline comments so it + can be inspected and manually edited by a support engineer. Example:: + + [general] + # Data sharing consent (true / false). + # consent = true + # + # Date until which consent is active (ISO). Auto-expires to false. + # consent_date = + """ + file_name: str = "datacollector.config" + + def __init__(self) -> None: + self.file_path = Path(_CONF_DIR) / self.file_name + self._cp = configparser.ConfigParser(interpolation=None) + self._lock = threading.Lock() + self._read() + + def _read(self) -> None: + """Read the config file if it exists; otherwise leave defaults.""" + if self.file_path.exists(): + self._cp.read(self.file_path) + else: + logging.info("No datacollector config found; using defaults.") + + def _write(self) -> None: + """Write the current config to file, creating parent directories if needed.""" + self.file_path.parent.mkdir(parents=True, exist_ok=True) + self._ensure_section("general") + + # consent + if self.consent is None: + try: + self._cp.remove_option("general", "consent") + except configparser.NoOptionError: + pass + else: + self._cp.set("general", "consent", "true" if self.consent else "false") + + # consent_date (stored as local date: YYYY-MM-DD) + consent_day = self.consent_date + if consent_day is None: + try: + self._cp.remove_option("general", CONSENT_DATE_KEY) + except configparser.NoOptionError: + pass + else: + self._cp.set("general", CONSENT_DATE_KEY, consent_day.isoformat()) + + with self.file_path.open("w", encoding="utf-8") as fh: + self._cp.write(fh) + + os.chmod(str(self.file_path), 0o600) + + def _ensure_section(self, section: str) -> None: + """ + Ensure that section exists in the config, creating it if necessary. + :param section: Section name to ensure. + """ + if not self._cp.has_section(section): + self._cp.add_section(section) + + @property + def consent(self) -> Optional[bool]: + """ + Get the consent state, or None if not yet set. + :return: True if consented, False if declined, None if undecided. + """ + try: + return self._cp.getboolean("general", "consent") + except (configparser.NoSectionError, configparser.NoOptionError): + return None + except ValueError: + return None + + @consent.setter + def consent(self, value: bool) -> None: + """ + Set consent to true or false, and clear consent expiry. If none, the consent is cleared + and becomes undecided. + :param value: True to opt in, False to opt out. + """ + with self._lock: + self._ensure_section("general") + if value is None: + self.clear_consent() + else: + self._cp.set("general", "consent", "true" if value else "false") + self._cp.remove_option("general", CONSENT_DATE_KEY) + self._write() + + def clear_consent(self) -> None: + """ + Unset consent so it becomes undecided again. + :return: None + """ + with self._lock: + self._ensure_section("general") + self._cp.remove_option("general", "consent") + self._cp.remove_option("general", CONSENT_DATE_KEY) + self._write() + + @property + def consent_date(self) -> Optional[date]: + """ + Return the consent expiry date as a local-date value, or None when unset. + When this date is reached, consent automatically expires to False. + :return: Local date of consent expiry, or None if not set. + """ + try: + value = self._cp.get("general", CONSENT_DATE_KEY) + except (configparser.NoSectionError, configparser.NoOptionError): + return None + value = value.strip() + if not value: + return None + # Preferred format: YYYY-MM-DD (local date only). + try: + return date.fromisoformat(value) + except ValueError: + pass + + return None + + def set_consent_with_expiry(self, consent_date: date) -> None: + """ + Enable consent until consent_date (inclusive), then auto-expire to False. + + When consent_date is today or tomorrow, record() uses 100% collection + probability; otherwise the default 10% applies. + + :param consent_date: Local date after which consent is revoked automatically. + :return: None + """ + with self._lock: + self._ensure_section("general") + self._cp.set("general", "consent", "true") + self._cp.set("general", CONSENT_DATE_KEY, consent_date.isoformat()) + self._write() + + def get_upload_backend(self) -> "S3UploadBackend": + """ + Return the configured upload backend instance. + + When the environment variable TEST_DATACOLLECTION is set to "1", + uploads are redirected to the test bucket (S3_TEST_BUCKET) so that + developer and CI runs do not pollute the production dataset. + + :return: Configured S3UploadBackend instance. + """ + credentials = _search_credentials() + if os.environ.get("TEST_DATACOLLECTION") == "1": + bucket = S3_TEST_BUCKET + logging.info( + "DataCollector: TEST_DATACOLLECTION=1 — using test bucket '%s'", bucket + ) + else: + bucket = S3_BUCKET + return S3UploadBackend( + access_key=credentials["access_key"], + secret_key=credentials["secret_key"], + endpoint_url=S3_ENDPOINT_URL, + region=S3_REGION, + bucket=bucket, + ) + + +@dataclass +class _WorkItem: + """A single data-collection event to be serialised and uploaded.""" + event_name: str + schema_version: str + payload: dict + image_format: str = "TIFF" + submitted_at: float = field(default_factory=time.monotonic) + + +def _serialize(item: _WorkItem, queue_dir: Path) -> Path: + """ + Serialise item into a ZIP archive and place it in queue_dir. + :param item: The work item to serialise. + :param queue_dir: Directory where the finished ZIP is placed. + :returns: Path to the created ZIP file inside queue_dir. + :raises OSError: On disk errors (caller must handle). + """ + queue_dir.mkdir(parents=True, exist_ok=True) + + sample_uuid = str(uuid.uuid4()) + uuid8 = sample_uuid.split("-")[0] + timestamp_utc = datetime.now(timezone.utc) + timestamp_str = timestamp_utc.strftime("%Y%m%dT%H%M%S") + # Sanitize and truncate event_name so the filename is always filesystem-safe. + safe_event = _sanitize_filename(item.event_name)[:64] if item.event_name else "event" + zip_name = f"{safe_event}-{timestamp_str}-{uuid8}.zip" + + tmp_dir = Path(tempfile.mkdtemp(prefix="dc_")) + try: + payload_meta: dict = {} + extra_files: list = [] # list of (arcname, abs_path) + + for key, value in item.payload.items(): + abs_path = None + if value is None or isinstance(value, (str, int, float, bool)): + payload_meta[key] = value + + elif isinstance(value, numpy.ndarray): + if item.image_format.upper() == "HDF5": + exporter = hdf5 + elif item.image_format.upper() == "TIFF": + exporter = tiff + else: + logging.warning("DataArray not in valid format", exc_info=True) + exporter = None + + if exporter is not None: + ext = exporter.EXTENSIONS[0] + arc_name = f"{_sanitize_filename(key)}.{ext}" + abs_path = tmp_dir / arc_name + try: + da = value if isinstance(value, model.DataArray) else model.DataArray(value) + tiff.export(str(abs_path), da) + except Exception: + logging.warning("Failed to export DataArray to %s at %s", ext, abs_path, exc_info=True) + abs_path = None + + if abs_path is not None and abs_path.exists(): + extra_files.append((arc_name, abs_path)) + payload_meta[key] = arc_name + else: + payload_meta[key] = None + payload_meta["export_error"] = True + + elif isinstance(value, (dict, list)): + arc_name = f"extra_{_sanitize_filename(key)}.json" + abs_path = tmp_dir / arc_name + abs_path.write_text(json.dumps(value, default=str), encoding="utf-8") + extra_files.append((arc_name, abs_path)) + payload_meta[key] = arc_name + + else: + # Fallback: store string representation, guarding against + # __repr__/__str__ implementations that raise. + try: + payload_meta[key] = str(value) + except Exception: + logging.warning("Failed to convert payload key '%s' to string", key, exc_info=True) + payload_meta[key] = "" + + metadata = { + "sample_uuid": sample_uuid, + "timestamp_utc": timestamp_utc.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z", + "system_id": socket.gethostname(), + "odemis_version": odemis.__version__, + "event_name": item.event_name, + "schema_version": item.schema_version, + "payload": payload_meta, + } + + meta_path = tmp_dir / "metadata.json" + meta_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") + + # Build ZIP in temp dir, then rename atomically into queue_dir. + tmp_zip = queue_dir / f"{uuid8}.tmp" + final_zip = queue_dir / zip_name + with zipfile.ZipFile(str(tmp_zip), "w", compression=zipfile.ZIP_DEFLATED) as zf: + zf.write(str(meta_path), "metadata.json") + for arc_name, abs_path in extra_files: + zf.write(str(abs_path), arc_name) + + os.replace(str(tmp_zip), str(final_zip)) + return final_zip + + finally: + shutil.rmtree(str(tmp_dir), ignore_errors=True) + + +def _enforce_queue_limit(queue_dir: Path) -> None: + """ + Delete the oldest ZIP files if the queue exceeds 10% of partition space. + :param queue_dir: The staging directory to inspect. + """ + if not queue_dir.exists(): + return + + zips = sorted(queue_dir.glob("*.zip"), key=lambda p: p.stat().st_mtime) + if not zips: + return + + try: + usage = shutil.disk_usage(str(queue_dir)) + except OSError: + logging.warning("Cannot read disk usage for %s", queue_dir) + return + + limit = usage.total * 0.10 # 10 % of partition + total_size = sum(p.stat().st_size for p in zips) + + while total_size > limit and zips: + oldest = zips.pop(0) + try: + size = oldest.stat().st_size + oldest.unlink() + total_size -= size + logging.info("Queue limit exceeded: removed oldest sample %s", oldest.name) + except OSError: + logging.warning("Could not remove queue file %s", oldest) + + +class S3UploadBackend: + """S3 upload backend implemented with boto3.""" + + def __init__( + self, + access_key: str, + secret_key: str, + endpoint_url: Optional[str] = S3_ENDPOINT_URL, + region: str = S3_REGION, + bucket: str = S3_BUCKET, + ) -> None: + """ + Initialize the S3 upload backend with the given credentials and configuration. + :param access_key: AWS access key ID. + :param secret_key: AWS secret access key. + :param endpoint_url: Optional S3 endpoint URL (for custom S3-compatible storage). + :param region: AWS region name (default "eu-west-1"). + :param bucket: S3 bucket name to upload to (default "delmic-odemis-collect"). + """ + self._access_key = access_key + self._secret_key = secret_key + self._endpoint_url = endpoint_url + self._region = region + self._bucket = bucket + self._client = None + + def _get_client(self) -> "boto3.client": + """ + Get a cached boto3 S3 client. + :return: boto3 S3 client instance. + """ + if self._client is None: + self._client = boto3.client( + "s3", + endpoint_url=self._endpoint_url, + region_name = self._region, + aws_access_key_id=self._access_key, + aws_secret_access_key=self._secret_key, + ) + return self._client + + def upload(self, local_path: Path, remote_key: str) -> None: + """Upload local_path to remote_key in the configured bucket.""" + client = self._get_client() + client.upload_file(str(local_path), self._bucket, remote_key) + + +def _upload(zip_path: Path, backend: S3UploadBackend) -> None: + """Upload zip_path with backend using the standard remote key.""" + remote_key = f"{socket.gethostname()}/{zip_path.name}" + backend.upload(zip_path, remote_key) + + +class _BackgroundWorker: + """ + Daemon thread that consumes WorkItem objects from a queue. + For each item it calls _enforce_queue_limit, _serialize, + and _upload in sequence. After a successful upload the local ZIP + is deleted. Exceptions are caught and logged, but never directly shown to the user. + The thread is started lazily and restarted automatically if it dies. + """ + + def __init__(self, config: DataCollectorConfig, queue_dir: Path = _DEFAULT_QUEUE_DIR) -> None: + self._config = config + self._queue_dir = queue_dir + self._queue: queue.Queue = queue.Queue() + self._thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + self._upload_backend: Optional[S3UploadBackend] = None + self._next_retry_at: float = 0.0 + self._retry_delay: float = _INITIAL_RETRY_DELAY_SECONDS + """ + Initialize the background worker with the given configuration and queue directory. + :param config: DataCollectorConfig instance for accessing configuration and upload backend. + :param queue_dir: Directory where ZIP files are staged for upload (default /var/log + /odemis/dc_queue). + """ + + def enqueue(self, item: _WorkItem) -> None: + """Add item to the processing queue and ensure the thread is alive. + :param item: The work item to enqueue. + """ + self._ensure_thread() + self._queue.put_nowait(item) + + def _ensure_thread(self) -> None: + """Start the background thread if it's not already running.""" + with self._lock: + if self._thread is None or not self._thread.is_alive(): + self._thread = threading.Thread( + target=self._run, + name="DataCollectorWorker", + daemon=True, + ) + self._thread.start() + logging.debug("DataCollector background thread started.") + + def _get_upload_backend(self) -> S3UploadBackend: + """Return a cached upload backend.""" + if self._upload_backend is None: + self._upload_backend = self._config.get_upload_backend() + return self._upload_backend + + def _schedule_retry(self) -> None: + """Schedule the next retry using exponential backoff.""" + delay = self._retry_delay + self._next_retry_at = time.monotonic() + delay + self._retry_delay = min(self._retry_delay * 2.0, _MAX_RETRY_DELAY_SECONDS) + logging.warning("DataCollector upload failed; retrying in %.0f s", delay) + + def _reset_retry(self) -> None: + """Reset retry state after a successful upload.""" + self._next_retry_at = 0.0 + self._retry_delay = _INITIAL_RETRY_DELAY_SECONDS + + def _pending_zip_paths(self, queue_dir: Path) -> list[Path]: + """Return pending ZIP files ordered oldest-first.""" + if not queue_dir.exists(): + return [] + return sorted(queue_dir.glob("*.zip"), key=lambda p: p.stat().st_mtime) + + def _process_pending_zips(self, queue_dir: Path) -> bool: + """ + Upload pending ZIP files from queue_dir. + :param queue_dir: Directory to scan for pending ZIP files. + :return: True when pending work existed (including backoff wait), + otherwise False. + """ + pending = self._pending_zip_paths(queue_dir) + if not pending: + return False + + now = time.monotonic() + if now < self._next_retry_at: + time.sleep(min(1.0, self._next_retry_at - now)) + return True + + try: + backend = self._get_upload_backend() + except Exception: + logging.warning("DataCollector failed to initialize upload backend", exc_info=True) + self._schedule_retry() + return True + for zip_path in pending: + try: + _upload(zip_path, backend) + zip_path.unlink(missing_ok=True) + self._reset_retry() + except Exception: + logging.warning( + "DataCollector upload failed for %s", zip_path.name, exc_info=True + ) + self._schedule_retry() + return True + return True + + def _process_work_item(self, item: _WorkItem) -> None: + """Serialize one work item and trigger upload processing.""" + _enforce_queue_limit(self._queue_dir) + _serialize(item, self._queue_dir) + self._process_pending_zips(self._queue_dir) + + def _run(self) -> None: + """Main loop: process items until the thread is stopped.""" + while True: + # Always drain the in-memory queue first with a non-blocking get so + # new record() calls are serialised to disk even when upload backoff + # is in progress. Without this, the queue grows unbounded and + # queue.get() is never reached while pending ZIPs exist. + try: + item = self._queue.get_nowait() + except queue.Empty: + item = None + + if item is not None: + try: + self._process_work_item(item) + except Exception: + logging.warning("DataCollector error processing event '%s'", + item.event_name, exc_info=True) + continue # immediately check for more queued items + + # No in-memory items; try to upload pending ZIPs from disk. + try: + had_pending = self._process_pending_zips(self._queue_dir) + except Exception: + logging.warning("DataCollector error while processing pending uploads", exc_info=True) + self._schedule_retry() + had_pending = True + + if had_pending: + continue + + # Nothing pending; block briefly for new items to arrive. + try: + item = self._queue.get(timeout=1.0) + except queue.Empty: + continue + try: + self._process_work_item(item) + except Exception: + logging.warning( + "DataCollector error processing event '%s'", item.event_name, exc_info=True + ) + + +class DataCollector: + """Thread-safe recorder for annotated data samples.""" + + def __init__(self) -> None: + self._config: Optional[DataCollectorConfig] = None + self._worker: Optional[_BackgroundWorker] = None + self._init_ok: bool = False + self._init_lock = threading.Lock() + + def _lazy_init(self) -> None: + """Initialise configuration and worker on first use.""" + if self._init_ok: + return + with self._init_lock: + if self._init_ok: + return + try: + self._config = DataCollectorConfig() + self._worker = _BackgroundWorker(self._config) + self._init_ok = True + logging.debug("DataCollector initialised.") + except Exception: + logging.warning( + "DataCollector failed to initialise; all record() calls will be no-ops." + ) + + def get_consent(self) -> Optional[bool]: + """Return current consent state and auto-expire temporary consent when needed.""" + self._lazy_init() + if not self._init_ok: + return None + consent = self._config.consent + if consent is not True: + return consent + + consent_day = self._config.consent_date + if consent_day is None: + return True + + today_local = datetime.now().astimezone().date() + if today_local > consent_day: + self._config.consent = False + logging.info("DataCollector: temporary consent expired; consent set to False.") + return False + return True + + def set_consent(self, value: bool) -> None: + """ + Persist explicit user consent choice. + :param value: Boolean indicating user's consent choice. + """ + if not isinstance(value, bool): + raise ValueError("value must be a bool") + self._lazy_init() + if not self._init_ok: + return + self._config.consent = value + + def set_temporary_consent(self, days: int = 1) -> None: + """ + Enable consent for the specified number of days, after which it + auto-expires to False. + + When days == 1, the remaining time at the moment record() is called + will be ≤ 1 day, so 100% collection probability applies throughout. + When days > 1, the default 10% probability applies until the final day. + + :param days: Number of days the temporary consent is active. + """ + self._lazy_init() + if not self._init_ok: + return + consent_day = datetime.now().astimezone().date() + timedelta(days=days) + self._config.set_consent_with_expiry(consent_date=consent_day) + + def record( + self, + event_name: str, + schema_version: str, + payload: dict, + image_format: str = "TIFF", + ) -> None: + """ + Capture an annotated data sample at a software event. + Returns immediately (non-blocking). Serialisation and upload happen + asynchronously in a background thread. If consent has not been + granted, this is a no-op. This function never raises (beyond the + input validation below); all errors are logged and suppressed. + + Collection probability is applied before enqueuing: 100% when + temporary consent (1-day) is active, 10% otherwise. + + :param event_name: Human-readable event identifier, e.g. + "z_stack_acquired". Must be a non-empty string. + :param schema_version: Payload schema version string, e.g. "1.0". + Must be a non-empty string. + :param payload: Dict of arbitrary values. Must be a dict. Supported + value types: + - Python primitives (str, int, float, bool, None) — inlined in + metadata.json + - :class:odemis.model.DataArray / :class:numpy.ndarray — + exported as TIFF or HDF5 files + - dict / list — written as extra_.json + :param image_format: Format for DataArray export. "TIFF" + (default) or "HDF5". + :raises ValueError: If any input parameter is invalid. + """ + if not isinstance(event_name, str) or not event_name: + raise ValueError("event_name must be a non-empty string") + if not isinstance(schema_version, str) or not schema_version: + raise ValueError("schema_version must be a non-empty string") + if not isinstance(payload, dict): + raise ValueError("payload must be a dict") + if not isinstance(image_format, str) or image_format.upper() not in _VALID_IMAGE_FORMATS: + raise ValueError( + f"image_format must be one of {_VALID_IMAGE_FORMATS}, got {image_format!r}" + ) + + try: + self._lazy_init() + if not self._init_ok: + return + + consent = self.get_consent() + if not consent: + logging.debug( + "DataCollector: consent=%s, skipping event '%s'.", consent, event_name + ) + return + + # Apply collection probability: 100% if consent expires within 1 day, + # 10% otherwise (no consent_date = permanent opt-in). + consent_day = self._config.consent_date + if consent_day is not None: + days_left = (consent_day - datetime.now().astimezone().date()).days + else: + days_left = None + + if days_left is not None and days_left <= 1: + probability = _FULL_COLLECTION_PROBABILITY + else: + probability = _DEFAULT_COLLECTION_PROBABILITY + + if random.random() >= probability: + logging.debug( + "DataCollector: event '%s' not sampled (%.0f%% collection probability).", + event_name, probability * 100, + ) + return + + item = _WorkItem( + event_name=event_name, + schema_version=schema_version, + payload=payload, + image_format=image_format, + ) + self._worker.enqueue(item) + except Exception: + logging.warning( + "Unexpected error in DataCollector.record(); event '%s' dropped.", event_name, exc_info=True + ) diff --git a/src/odemis/util/dc_fetch.py b/src/odemis/util/dc_fetch.py new file mode 100644 index 0000000000..d620c3dc0d --- /dev/null +++ b/src/odemis/util/dc_fetch.py @@ -0,0 +1,336 @@ +# -*- coding: utf-8 -*- +""" +Created on 11 March 2026 + +@author: Karishma Kumar + +Copyright © 2026 Karishma Kumar, Delmic + +This file is part of Odemis. + +Odemis is free software: you can redistribute it and/or modify it under the +terms of the GNU General Public License version 2 as published by the Free +Software Foundation. + +Odemis is distributed in the hope that it will be useful, but WITHOUT ANY +WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR +A PARTICULAR PURPOSE. See the GNU General Public License for more details. + +You should have received a copy of the GNU General Public License along with +Odemis. If not, see http://www.gnu.org/licenses/. + +Retrieval helpers for downloading DataCollector ZIP samples from S3. +""" + +import logging +import configparser +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Tuple + +import boto3 + +from odemis.util.datacollector import S3_BUCKET, S3_ENDPOINT_URL, S3_REGION + + +_DC_FETCH_CONFIG_PATH = Path.home() / ".config" / "odemis" / "dc_fetch.ini" + + +def parse_since_utc(value: str) -> datetime: + """ + Parse a date/datetime string to UTC-aware datetime. + :param value: ISO-8601 date (YYYY-MM-DD) or datetime + (YYYY-MM-DDTHH:MM:SS with optional timezone offset or Z suffix). + :return: UTC-aware datetime. + """ + text = value.strip() + if len(text) == 10: + parsed = datetime.strptime(text, "%Y-%m-%d") + return parsed.replace(tzinfo=timezone.utc) + # datetime.fromisoformat() does not accept the 'Z' suffix in Python < 3.11. + if text.endswith("Z"): + text = text[:-1] + "+00:00" + parsed = datetime.fromisoformat(text) + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + +def parse_key_timestamp_utc(key: str) -> Optional[datetime]: + """ + Parse --.zip timestamp from key basename. + + The S3 object key is the full path-like key in the bucket, for example: + ``meteor-5099/z_stack_acquired-20260322T104530-a1b2c3d4.zip``. + In this format, ``meteor-5099/`` is the host prefix and the basename is + ``z_stack_acquired-20260322T104530-a1b2c3d4.zip``. + + :param key: S3 object key. + :return: Parsed UTC datetime, or None if parsing failed. + """ + name = Path(key).name + if not name.endswith(".zip"): + return None + stem = name[:-4] + parts = stem.rsplit("-", 2) + if len(parts) != 3: + return None + ts = parts[1] + try: + parsed = datetime.strptime(ts, "%Y%m%dT%H%M%S") + except ValueError: + return None + return parsed.replace(tzinfo=timezone.utc) + +def parse_key_event_name(key: str) -> Optional[str]: + """ + Parse event name from --.zip key basename. + + Example key: + ``meteor-5099/z_stack_acquired-20260322T104530-a1b2c3d4.zip`` + + Parsed event name from the basename: + ``z_stack_acquired`` + + :param key: S3 object key. + :return: Event name, or None if parsing failed. + """ + name = Path(key).name + if not name.endswith(".zip"): + return None + stem = name[:-4] + parts = stem.rsplit("-", 2) + if len(parts) != 3: + return None + return parts[0] or None + +def iter_s3_objects(s3_client: Any, bucket: str, prefix: str) -> Iterator[Dict[str, Any]]: + """ + Iterate S3 objects under prefix using list_objects_v2 pagination. + :param s3_client: Boto3 S3 client instance. + :param bucket: S3 bucket name. + :param prefix: S3 prefix to filter objects. + :return: Iterator of S3 object dictionaries. + """ + token: Optional[str] = None + while True: + kwargs: Dict[str, Any] = {"Bucket": bucket, "Prefix": prefix} + if token: + kwargs["ContinuationToken"] = token + response = s3_client.list_objects_v2(**kwargs) + for item in response.get("Contents", []): + yield item + if not response.get("IsTruncated"): + break + token = response.get("NextContinuationToken") + if not token: + logging.warning("S3 list_objects_v2 returned IsTruncated=True but no NextContinuationToken; stopping pagination.") + break + +def should_download_key(key: str, event_filter: Optional[str], since_utc: Optional[datetime]) -> bool: + """ + Return whether an S3 key should be downloaded by filters. + :param key: S3 object key (for example: ``meteor-5099/z_stack_acquired-20260322T104530-a1b2c3d4.zip``). + :param event_filter: Optional event name filter. + :param since_utc: Optional UTC datetime filter. + :return: True if the key should be downloaded, False otherwise. + """ + if not key.endswith(".zip"): + return False + if event_filter: + event_name = parse_key_event_name(key) + if event_name != event_filter: + return False + if since_utc: + key_ts = parse_key_timestamp_utc(key) + if key_ts is None: + return False + if key_ts < since_utc: + return False + return True + +def parse_host_filters(value: Optional[str]) -> List[str]: + """ + Parse comma-separated host filters into normalized host IDs. + :param value: Comma-separated host filter string. + :return: List of normalized host IDs. + """ + if not value: + return [] + hosts = [part.strip().strip("/") for part in value.split(",")] + return [host for host in hosts if host] + + +def _write_dc_fetch_config( + config: configparser.ConfigParser, + config_path: Path, +) -> None: + """ + Write dc_fetch INI config to disk, creating parent directories if needed. + + :param config: Parsed configuration object. + :param config_path: Destination INI file path. + :return: None. + """ + config_path.parent.mkdir(parents=True, exist_ok=True) + with config_path.open("w", encoding="utf-8") as fp: + config.write(fp) + + +def _load_or_init_dc_fetch_config(config_path: Path) -> configparser.ConfigParser: + """ + Load the data retrieval keys from the config path. If the file does not + exist, the template for the file will be created. + + :param config_path: path from where the data retrieval keys exists + :return: Parsed configuration object. + :raises RuntimeError: If the config was just created or credentials are missing. + """ + resolved_path = config_path or _DC_FETCH_CONFIG_PATH + cp = configparser.ConfigParser(interpolation=None) + default_endpoint = "" if S3_ENDPOINT_URL is None else str(S3_ENDPOINT_URL) + defaults = { + "access_key": "", + "secret_key": "", + "bucket": S3_BUCKET, + "endpoint_url": default_endpoint, + "region": S3_REGION, + } + + if resolved_path.exists(): + cp.read(str(resolved_path), encoding="utf-8") + else: + cp["s3"] = defaults + _write_dc_fetch_config(cp, resolved_path) + raise RuntimeError( + "dc_fetch config created at {}. Please edit [s3] access_key and " + "secret_key, then run the command again.".format(resolved_path) + ) + + if not cp.has_section("s3"): + cp["s3"] = defaults + _write_dc_fetch_config(cp, resolved_path) + raise RuntimeError( + "Missing [s3] section in {}. A template was written. Please set " + "[s3] access_key and secret_key.".format(resolved_path) + ) + + access_key = cp.get("s3", "access_key", fallback="").strip() + secret_key = cp.get("s3", "secret_key", fallback="").strip() + if not access_key or not secret_key: + raise RuntimeError( + "Missing S3 credentials in {}. Please set [s3] access_key and " + "secret_key.".format(resolved_path) + ) + + return cp + +def build_s3_client_from_config( + config_path: Path, + bucket_override: Optional[str] = None, + endpoint_override: Optional[str] = None, + region_override: Optional[str] = None, +) -> Tuple[Any, str]: + """ + Build an S3 client from dc_fetch.ini with optional endpoint/bucket overrides. + :param config_path: path to retrieve keys for data retrieval. + :param bucket_override: Optional S3 bucket name override. + :param endpoint_override: Optional S3 endpoint URL override. + :param region_override: Optional AWS region name override. + :return: Tuple of (Boto3 S3 client, bucket name). + """ + cp = _load_or_init_dc_fetch_config(config_path=config_path) + default_endpoint = "" if S3_ENDPOINT_URL is None else str(S3_ENDPOINT_URL) + + endpoint_url_text = endpoint_override + if endpoint_url_text is None: + endpoint_url_text = cp.get("s3", "endpoint_url", fallback=default_endpoint).strip() + endpoint_url = endpoint_url_text or None + + bucket = bucket_override + if bucket is None: + bucket = cp.get("s3", "bucket", fallback=S3_BUCKET).strip() or S3_BUCKET + + region_name = region_override + if region_name is None: + region_name = cp.get("s3", "region", fallback=S3_REGION).strip() or S3_REGION + + client_kwargs: Dict[str, Any] = { + "endpoint_url": endpoint_url, + "aws_access_key_id": cp.get("s3", "access_key").strip(), + "aws_secret_access_key": cp.get("s3", "secret_key").strip(), + "region_name": region_name, + } + client = boto3.client("s3", **client_kwargs) + return client, bucket + +def fetch_samples( + event_filter: Optional[str], + since_utc: Optional[datetime], + output_dir: Path, + host_filter: Optional[str] = None, + bucket_override: Optional[str] = None, + endpoint_override: Optional[str] = None, + region_override: Optional[str] = None, +) -> Dict[str, int]: + """ + Fetch matching samples from S3 into output directory. + :param event_filter: Optional event name filter. + :param since_utc: Optional UTC datetime filter. + :param output_dir: Directory to save downloaded samples. + :param host_filter: Optional comma-separated host filter string. + :param bucket_override: Optional S3 bucket name override. + :param endpoint_override: Optional S3 endpoint URL override. + :param region_override: Optional AWS region name override. + :return: Dictionary with counts of listed, matched, downloaded, skipped, and failed samples. + """ + s3_client, bucket = build_s3_client_from_config( + config_path=_DC_FETCH_CONFIG_PATH, + bucket_override=bucket_override, + endpoint_override=endpoint_override, + region_override=region_override, + ) + host_filters = parse_host_filters(host_filter) + prefixes = [f"{host}/" for host in host_filters] if host_filters else [""] + + output_dir.mkdir(parents=True, exist_ok=True) + + listed = 0 + matched = 0 + downloaded = 0 + skipped_existing = 0 + failed = 0 + + for prefix in prefixes: + for item in iter_s3_objects(s3_client, bucket=bucket, prefix=prefix): + listed += 1 + raw_key = item.get("Key") + key = raw_key if isinstance(raw_key, str) else None + if not key or not should_download_key(key, event_filter, since_utc): + continue + matched += 1 + # Flatten the S3 key (e.g. "host/file.zip" → "host_file.zip") so + # files from different hosts never collide in the output directory. + flat_name = key.replace("/", "_") + destination = output_dir / flat_name + if destination.exists(): + skipped_existing += 1 + continue + # Write to a .part file first so a failed download never leaves a + # truncated ZIP that would be mistaken for a complete file on retry. + tmp_dest = Path(str(destination.with_suffix(".part"))) + try: + s3_client.download_file(bucket, key, str(tmp_dest)) + tmp_dest.rename(destination) + downloaded += 1 + except Exception: + tmp_dest.unlink(missing_ok=True) + failed += 1 + logging.exception("Failed to download key %s", key) + + return { + "listed": listed, + "matched": matched, + "downloaded": downloaded, + "skipped_existing": skipped_existing, + "failed": failed, + } diff --git a/src/odemis/util/test/datacollector_test.py b/src/odemis/util/test/datacollector_test.py new file mode 100644 index 0000000000..1227a34a93 --- /dev/null +++ b/src/odemis/util/test/datacollector_test.py @@ -0,0 +1,723 @@ +# -*- coding: utf-8 -*- +""" +Created on 11 March 2026 + +@author: Karishma Kumar + +Copyright © 2026 Delmic + +This file is part of Odemis. + +Odemis is free software: you can redistribute it and/or modify it under the +terms of the GNU General Public License version 2 as published by the Free +Software Foundation. + +Odemis is distributed in the hope that it will be useful, but WITHOUT ANY +WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A +PARTICULAR PURPOSE. See the GNU General Public License for more details. + +You should have received a copy of the GNU General Public License along with +Odemis. If not, see http://www.gnu.org/licenses/. +""" + +import configparser +import json +import logging +import os +import shutil +import socket +import stat +import tempfile +import time +import unittest +import uuid +import zipfile +from datetime import date, datetime, timedelta +from pathlib import Path +from unittest.mock import patch + +import numpy + +from odemis.util.datacollector import ( + DataCollector, + DataCollectorConfig, + S3UploadBackend, + S3_REGION, + S3_TEST_BUCKET, + _CREDENTIALS_PATH, + _BackgroundWorker, + _WorkItem, + _enforce_queue_limit, + _serialize, +) + + +class TestDataCollectorConfig(unittest.TestCase): + """Tests for DataCollectorConfig read/write behaviour.""" + + def setUp(self) -> None: + self._tmp_conf_dir = tempfile.mkdtemp(prefix="dc_conf_") + + def tearDown(self) -> None: + shutil.rmtree(self._tmp_conf_dir, ignore_errors=True) + + def _make_config(self) -> DataCollectorConfig: + """Return a DataCollectorConfig pointed at the temp directory.""" + cfg = DataCollectorConfig.__new__(DataCollectorConfig) + cfg.file_path = Path(self._tmp_conf_dir) / "datacollector.config" + import threading + cfg._cp = configparser.ConfigParser(interpolation=None) + cfg._lock = threading.Lock() + cfg._read() + return cfg + + def test_consent_round_trip(self) -> None: + """Setting consent to True/False persists to disk and re-reads correctly.""" + cfg = self._make_config() + cfg.consent = True + cfg2 = self._make_config() + self.assertTrue(cfg2.consent) + cfg.consent = False + cfg3 = self._make_config() + self.assertFalse(cfg3.consent) + + def test_clear_consent_round_trip(self) -> None: + """Cleared consent writes a file that reloads as None (not a ValueError).""" + cfg = self._make_config() + cfg.consent = True + cfg.clear_consent() + cfg2 = self._make_config() + self.assertIsNone(cfg2.consent) + + def test_legacy_consent_none_string_reads_as_none(self) -> None: + """A config with 'consent = none' (legacy format) must read as None.""" + cfg = self._make_config() + cfg.file_path.parent.mkdir(parents=True, exist_ok=True) + cfg.file_path.write_text("[general]\nconsent = none\n", encoding="utf-8") + cfg._cp.read(str(cfg.file_path)) + self.assertIsNone(cfg.consent) + + def test_config_file_permissions(self) -> None: + """Config file should be written with mode 0o600 (security requirement).""" + cfg = self._make_config() + cfg.consent = True + mode = stat.S_IMODE(os.stat(str(cfg.file_path)).st_mode) + self.assertEqual(mode, 0o600) + + def test_set_consent_with_expiry_stores_fields(self) -> None: + """set_consent_with_expiry should enable consent and set consent_date.""" + cfg = self._make_config() + expiry = date.today() + timedelta(days=1) + cfg.set_consent_with_expiry(expiry) + + self.assertTrue(cfg.consent) + stored_date = cfg.consent_date + self.assertIsNotNone(stored_date) + self.assertEqual(stored_date, expiry) + + +class TestTemporaryConsentAndProbability(unittest.TestCase): + """Tests for set_temporary_consent and record() expiry/probability logic.""" + + def setUp(self) -> None: + self._tmp_dir = Path(tempfile.mkdtemp(prefix="dc_tmp_consent_")) + self._queue_dir = self._tmp_dir / "queue" + self._queue_dir.mkdir(parents=True, exist_ok=True) + + cfg = DataCollectorConfig.__new__(DataCollectorConfig) + cfg.file_path = self._tmp_dir / "datacollector.config" + import threading + cfg._cp = configparser.ConfigParser(interpolation=None) + cfg._lock = threading.Lock() + cfg._read() + self._cfg = cfg + + self._dc = DataCollector.__new__(DataCollector) + self._dc._config = cfg + self._dc._worker = _BackgroundWorker(cfg, queue_dir=self._queue_dir) + self._dc._init_ok = True + import threading as _threading + self._dc._init_lock = _threading.Lock() + + def tearDown(self) -> None: + shutil.rmtree(str(self._tmp_dir), ignore_errors=True) + + def test_set_temporary_consent_enables_consent_with_expiry(self) -> None: + """set_temporary_consent(1) should set consent=True and consent_date to next local day.""" + expected_day = datetime.now().astimezone().date() + timedelta(days=1) + self._dc.set_temporary_consent(days=1) + + self.assertTrue(self._cfg.consent) + stored = self._cfg.consent_date + self.assertIsNotNone(stored) + # Midnight roll-over during the call is possible; allow +1 day tolerance. + self.assertIn(stored, {expected_day, expected_day + timedelta(days=1)}) + + def test_set_temporary_consent_1day_uses_full_probability(self) -> None: + """record() should enqueue all events when consent_date is ~1 day away.""" + self._dc.set_temporary_consent(days=1) + + enqueue_calls = [] + self._dc._worker.enqueue = lambda item: enqueue_calls.append(item) + + trials = 20 + for i in range(trials): + self._dc.record("full_event", "1.0", {"i": i}) + + self.assertEqual(len(enqueue_calls), trials, + "All events should be enqueued when consent expires within 1 day") + + def test_set_temporary_consent_multiday_uses_default_probability(self) -> None: + """record() should sample ~10% when consent_date is more than 1 day away.""" + future = date.today() + timedelta(days=30) + self._cfg.set_consent_with_expiry(future) + + enqueue_calls = [] + self._dc._worker.enqueue = lambda item: enqueue_calls.append(item) + + trials = 5000 + for i in range(trials): + self._dc.record("sampled_event", "1.0", {"i": i}) + + ratio = len(enqueue_calls) / trials + # Allow ±5% around the 10% target. + self.assertGreater(ratio, 0.05, f"Sampling ratio {ratio:.2%} too low") + self.assertLess(ratio, 0.15, f"Sampling ratio {ratio:.2%} too high") + + def test_record_auto_expires_consent_and_skips(self) -> None: + """record() should set consent=False and skip when consent_date has passed.""" + expired = datetime.now().astimezone().date() - timedelta(days=1) + self._cfg.set_consent_with_expiry(expired) + + enqueue_calls = [] + self._dc._worker.enqueue = lambda item: enqueue_calls.append(item) + + self._dc.record("expired_event", "1.0", {"x": 1}) + + self.assertFalse(self._cfg.consent) + self.assertEqual(len(enqueue_calls), 0, "record() should not enqueue after expiry") + + def test_record_no_consent_is_noop(self) -> None: + """record() should be a no-op when consent is False.""" + self._cfg.consent = False + + enqueue_calls = [] + self._dc._worker.enqueue = lambda item: enqueue_calls.append(item) + + self._dc.record("no_consent_event", "1.0", {"x": 1}) + self.assertEqual(len(enqueue_calls), 0) + + +class TestSerialize(unittest.TestCase): + """Tests for _serialize() — ZIP structure and metadata correctness.""" + + def setUp(self) -> None: + self._tmp_queue = Path(tempfile.mkdtemp(prefix="dc_queue_")) + + def tearDown(self) -> None: + shutil.rmtree(str(self._tmp_queue), ignore_errors=True) + + def _make_item(self, payload: dict, image_format: str = "TIFF") -> _WorkItem: + return _WorkItem( + event_name="test_event", + schema_version="1.0", + payload=payload, + image_format=image_format, + ) + + def test_zip_created(self) -> None: + """A ZIP file is created in queue_dir after serialisation.""" + item = self._make_item({"score": 0.9}) + zip_path = _serialize(item, self._tmp_queue) + self.assertTrue(zip_path.exists(), "ZIP file not created") + self.assertTrue(zip_path.suffix == ".zip") + + def test_zip_filename_format(self) -> None: + """ZIP filename follows --.zip convention.""" + item = self._make_item({"x": 1}) + zip_path = _serialize(item, self._tmp_queue) + name = zip_path.name + parts = name[:-4].split("-") # strip .zip + self.assertEqual(parts[0], "test_event") + self.assertEqual(len(parts[2]), 8, "UUID8 part should be 8 hex characters") + + def test_metadata_json_envelope_fields(self) -> None: + """metadata.json must contain all standard envelope fields.""" + item = self._make_item({"score": 0.5}) + zip_path = _serialize(item, self._tmp_queue) + with zipfile.ZipFile(str(zip_path)) as zf: + meta = json.loads(zf.read("metadata.json")) + required = {"sample_uuid", "timestamp_utc", "system_id", "odemis_version", + "event_name", "schema_version", "payload"} + self.assertEqual(required, required & meta.keys()) + self.assertEqual(meta["event_name"], "test_event") + self.assertEqual(meta["schema_version"], "1.0") + + def test_primitive_payload_inlined(self) -> None: + """Primitive payload values are inlined in metadata.json.""" + item = self._make_item({"score": 0.87, "n": 12, "name": "foo", "flag": True}) + zip_path = _serialize(item, self._tmp_queue) + with zipfile.ZipFile(str(zip_path)) as zf: + meta = json.loads(zf.read("metadata.json")) + self.assertAlmostEqual(meta["payload"]["score"], 0.87) + self.assertEqual(meta["payload"]["n"], 12) + self.assertEqual(meta["payload"]["name"], "foo") + self.assertTrue(meta["payload"]["flag"]) + + def test_numpy_array_exported_as_tiff(self) -> None: + """numpy.ndarray values are exported as .ome.tiff formats.""" + arr = numpy.zeros((64, 64), dtype=numpy.uint16) + item = self._make_item({"image": arr}) + zip_path = _serialize(item, self._tmp_queue) + with zipfile.ZipFile(str(zip_path)) as zf: + names = zf.namelist() + meta = json.loads(zf.read("metadata.json")) + output_formats = meta["payload"]["image"] + self.assertIn(output_formats, names, "TIFF sidecar not in ZIP") + self.assertTrue(output_formats.endswith(".ome.tiff")) + + def test_dict_payload_written_as_extra_json(self) -> None: + """dict payload values are written as extra_*.json.""" + item = self._make_item({"params": {"a": 1, "b": 2}}) + zip_path = _serialize(item, self._tmp_queue) + with zipfile.ZipFile(str(zip_path)) as zf: + names = zf.namelist() + meta = json.loads(zf.read("metadata.json")) + self.assertIn("extra_params.json", names) + self.assertEqual(meta["payload"]["params"], "extra_params.json") + + def test_list_payload_written_as_extra_json(self) -> None: + """list payload values are written as extra_*.json sidecars.""" + item = self._make_item({"items": [1, 2, 3]}) + zip_path = _serialize(item, self._tmp_queue) + with zipfile.ZipFile(str(zip_path)) as zf: + names = zf.namelist() + self.assertIn("extra_items.json", names) + + def test_atomic_write(self) -> None: + """No .tmp files remain after successful serialisation.""" + item = self._make_item({"x": 1}) + _serialize(item, self._tmp_queue) + tmps = list(self._tmp_queue.glob("*.tmp")) + self.assertEqual(tmps, [], "Leftover .tmp files found") + + def test_hdf5_image_format(self) -> None: + """When image_format=HDF5, DataArray is exported as an .h5.""" + arr = numpy.zeros((32, 32), dtype=numpy.float32) + item = self._make_item({"data": arr}, image_format="HDF5") + zip_path = _serialize(item, self._tmp_queue) + with zipfile.ZipFile(str(zip_path)) as zf: + names = zf.namelist() + meta = json.loads(zf.read("metadata.json")) + output_formats = meta["payload"]["data"] + self.assertIn(output_formats, names) + self.assertTrue(output_formats.endswith(".h5")) + + +class TestEnforceQueueLimit(unittest.TestCase): + """Tests for _enforce_queue_limit().""" + + def setUp(self) -> None: + self._tmp_queue = Path(tempfile.mkdtemp(prefix="dc_qlimit_")) + + def tearDown(self) -> None: + shutil.rmtree(str(self._tmp_queue), ignore_errors=True) + + def _write_zip(self, name: str, size_bytes: int, mtime: float) -> Path: + """Create a dummy ZIP file of the given size and modification time.""" + p = self._tmp_queue / name + p.write_bytes(b"\x00" * size_bytes) + os.utime(str(p), (mtime, mtime)) + return p + + def test_no_deletion_when_under_limit(self) -> None: + """Files are NOT deleted when total queue size is within the 10% limit.""" + self._write_zip("small.zip", 1024, time.time()) + _enforce_queue_limit(self._tmp_queue) + self.assertTrue((self._tmp_queue / "small.zip").exists()) + + def test_oldest_deleted_when_over_limit(self) -> None: + """Oldest ZIPs are deleted when the queue exceeds 10% of partition space.""" + import collections + FakeDiskUsage = collections.namedtuple("usage", ["total", "used", "free"]) + fake_usage = FakeDiskUsage(total=300, used=0, free=300) + + now = time.time() + # 3 files × 20 bytes = 60 bytes > 10% of 300 (= 30 bytes). + old = self._write_zip("old.zip", 20, now - 100) + self._write_zip("newer.zip", 20, now - 50) + self._write_zip("newest.zip", 20, now) + + with patch("odemis.util.datacollector.shutil.disk_usage", return_value=fake_usage): + _enforce_queue_limit(self._tmp_queue) + + self.assertFalse(old.exists(), "Oldest ZIP should have been deleted") + + def test_empty_dir_no_error(self) -> None: + """_enforce_queue_limit on an empty directory must not raise.""" + try: + _enforce_queue_limit(self._tmp_queue) + except Exception as exc: + self.fail(f"_enforce_queue_limit raised unexpectedly: {exc}") + + def test_nonexistent_dir_no_error(self) -> None: + """_enforce_queue_limit on a non-existent directory must not raise.""" + try: + _enforce_queue_limit(Path("/nonexistent/path/dc_queue")) + except Exception as exc: + self.fail(f"_enforce_queue_limit raised unexpectedly: {exc}") + + +class TestUploadAndRetry(unittest.TestCase): + """Tests upload backend and retry behavior.""" + + def setUp(self) -> None: + self._tmp_dir = Path(tempfile.mkdtemp(prefix="dc_upload_")) + self._queue_dir = self._tmp_dir / "queue" + self._queue_dir.mkdir(parents=True, exist_ok=True) + + cfg = DataCollectorConfig.__new__(DataCollectorConfig) + cfg.file_path = self._tmp_dir / "datacollector.config" + import threading + cfg._cp = configparser.ConfigParser(interpolation=None) + cfg._lock = threading.Lock() + cfg._read() + cfg.consent = True + self._cfg = cfg + self._worker = _BackgroundWorker(cfg, queue_dir=self._queue_dir) + + def tearDown(self) -> None: + shutil.rmtree(str(self._tmp_dir), ignore_errors=True) + + def _write_zip(self, name: str, mtime: float) -> Path: + """Create a pending ZIP in queue_dir with a stable mtime.""" + p = self._queue_dir / name + p.write_bytes(b"zip") + os.utime(str(p), (mtime, mtime)) + return p + + def test_upload_called_after_serialization(self) -> None: + """Worker should upload a serialized ZIP right after it is created.""" + order = [] + item = _WorkItem(event_name="upload_call_test", schema_version="1.0", payload={"x": 1}) + target_zip = self._queue_dir / "serialized.zip" + + class _Backend: + def upload(self, local_path: Path, remote_key: str) -> None: + order.append(("upload", local_path.name, remote_key)) + + def _fake_serialize(work_item: _WorkItem, queue_dir: Path) -> Path: + del work_item, queue_dir + target_zip.write_bytes(b"zip") + order.append(("serialize", target_zip.name)) + return target_zip + + with patch("odemis.util.datacollector._serialize", side_effect=_fake_serialize), \ + patch.object(self._worker, "_get_upload_backend", return_value=_Backend()): + self._worker._process_work_item(item) + + self.assertEqual(order[0], ("serialize", "serialized.zip")) + self.assertEqual(order[1][0], "upload") + self.assertEqual(order[1][1], "serialized.zip") + self.assertFalse(target_zip.exists(), "ZIP should be deleted after successful upload") + + def test_retry_on_failure_then_success(self) -> None: + """Failed upload is retried and eventually clears pending ZIPs.""" + now = time.time() + older = self._write_zip("older.zip", now - 60) + newer = self._write_zip("newer.zip", now - 30) + uploaded_names = [] + failures = {"count": 0} + + class _Backend: + def upload(self, local_path: Path, remote_key: str) -> None: + uploaded_names.append(local_path.name) + if failures["count"] == 0: + failures["count"] += 1 + raise ConnectionError("temporary network issue") + + with patch.object(self._worker, "_get_upload_backend", return_value=_Backend()): + had_pending = self._worker._process_pending_zips(self._queue_dir) + self.assertTrue(had_pending) + self.assertTrue(older.exists(), "Failed ZIP should remain for retry") + self.assertGreater(self._worker._next_retry_at, 0.0) + + with patch("odemis.util.datacollector.time.monotonic", return_value=self._worker._next_retry_at + 1.0): + had_pending = self._worker._process_pending_zips(self._queue_dir) + self.assertTrue(had_pending) + + self.assertFalse(older.exists()) + self.assertFalse(newer.exists()) + self.assertEqual(uploaded_names[:3], ["older.zip", "older.zip", "newer.zip"]) + + def test_pending_flush_oldest_first(self) -> None: + """Recovery flush should process queued ZIP files oldest-first.""" + now = time.time() + self._write_zip("oldest.zip", now - 120) + self._write_zip("middle.zip", now - 60) + self._write_zip("newest.zip", now - 10) + uploaded = [] + + class _Backend: + def upload(self, local_path: Path, remote_key: str) -> None: + uploaded.append(local_path.name) + + with patch.object(self._worker, "_get_upload_backend", return_value=_Backend()): + had_pending = self._worker._process_pending_zips(self._queue_dir) + self.assertTrue(had_pending) + self.assertEqual(uploaded, ["oldest.zip", "middle.zip", "newest.zip"]) + self.assertEqual(list(self._queue_dir.glob("*.zip")), []) + + +class TestRealS3Integration(unittest.TestCase): + """Real S3 integration tests for Phase 2 upload workflow. + + Tests use the credentials from _CREDENTIALS_PATH and upload to + S3_TEST_BUCKET (not the production bucket). The test class is + skipped automatically when the key file is absent. + """ + + @classmethod + def setUpClass(cls) -> None: + try: + import boto3 + except ImportError as exc: + raise unittest.SkipTest(f"boto3 is required for real S3 integration tests: {exc}") + + import json as _json + if not os.path.isfile(_CREDENTIALS_PATH): + raise unittest.SkipTest( + f"S3 key file not found at {_CREDENTIALS_PATH}; skipping real S3 integration tests." + ) + with open(_CREDENTIALS_PATH, "r") as fh: + creds = _json.load(fh) + + cls._access_key = creds["access_key"] + cls._secret_key = creds["secret_key"] + cls._bucket = S3_TEST_BUCKET + cls._region = S3_REGION + + cls._s3_client = boto3.client( + "s3", + aws_access_key_id=cls._access_key, + aws_secret_access_key=cls._secret_key, + region_name=cls._region, + ) + + def setUp(self) -> None: + self._tmp_dir = Path(tempfile.mkdtemp(prefix="dc_reals3_")) + self._queue_dir = self._tmp_dir / "queue" + self._queue_dir.mkdir(parents=True, exist_ok=True) + self._created_keys: list[str] = [] + + def tearDown(self) -> None: + for key in self._created_keys: + try: + self._s3_client.delete_object(Bucket=self._bucket, Key=key) + except Exception as exc: + logging.warning("Could not delete test object %s: %s", key, exc) + shutil.rmtree(str(self._tmp_dir), ignore_errors=True) + + def _new_remote_key(self, suffix: str = ".zip") -> str: + """Create a unique remote key for test uploads.""" + return f"odemis-integration-tests/{socket.gethostname()}/{uuid.uuid4().hex}{suffix}" + + def test_s3_upload_backend_uploads_file(self) -> None: + """S3UploadBackend.upload should place an object in the configured bucket.""" + local_path = self._tmp_dir / "sample.zip" + local_path.write_bytes(b"integration-test-payload") + remote_key = self._new_remote_key() + + backend = S3UploadBackend( + access_key=self._access_key, + secret_key=self._secret_key, + region=self._region, + bucket=self._bucket, + ) + backend.upload(local_path, remote_key) + self._created_keys.append(remote_key) + + response = self._s3_client.head_object(Bucket=self._bucket, Key=remote_key) + self.assertGreater(response["ContentLength"], 0) + + def test_worker_uploads_pending_and_deletes_local(self) -> None: + """Background worker should upload pending ZIPs and delete them locally.""" + old_zip = self._queue_dir / "old.zip" + new_zip = self._queue_dir / "new.zip" + old_zip.write_bytes(b"old") + new_zip.write_bytes(b"new") + now = time.time() + os.utime(str(old_zip), (now - 60, now - 60)) + os.utime(str(new_zip), (now - 30, now - 30)) + + cfg = DataCollectorConfig.__new__(DataCollectorConfig) + cfg.file_path = self._tmp_dir / "datacollector.config" + import threading + cfg._cp = configparser.ConfigParser(interpolation=None) + cfg._lock = threading.Lock() + cfg._read() + cfg.consent = True + worker = _BackgroundWorker(cfg, queue_dir=self._queue_dir) + + uploaded_keys: list[str] = [] + backend = S3UploadBackend( + access_key=self._access_key, + secret_key=self._secret_key, + region=self._region, + bucket=self._bucket, + ) + + def _capture_upload(local_path: Path, backend_obj: S3UploadBackend) -> None: + remote_key = f"{socket.gethostname()}/{local_path.name}" + backend_obj.upload(local_path, remote_key) + uploaded_keys.append(remote_key) + + with patch.object(worker, "_get_upload_backend", return_value=backend), \ + patch("odemis.util.datacollector._upload", side_effect=_capture_upload): + had_pending = worker._process_pending_zips(self._queue_dir) + + self.assertTrue(had_pending) + self.assertFalse(old_zip.exists(), "Old pending ZIP should be removed locally") + self.assertFalse(new_zip.exists(), "New pending ZIP should be removed locally") + self.assertEqual(len(uploaded_keys), 2) + + self._created_keys.extend(uploaded_keys) + for key in uploaded_keys: + self.assertIsNotNone(self._s3_client.head_object(Bucket=self._bucket, Key=key)) + + + +class DataCollectorTest(unittest.TestCase): + """Integration-level tests for DataCollector.record().""" + + def setUp(self) -> None: + self._tmp_dir = Path(tempfile.mkdtemp(prefix="dc_test_")) + self._queue_dir = self._tmp_dir / "queue" + self._queue_dir.mkdir(parents=True, exist_ok=True) + + cfg = DataCollectorConfig.__new__(DataCollectorConfig) + cfg.file_path = self._tmp_dir / "datacollector.config" + import threading + cfg._cp = configparser.ConfigParser(interpolation=None) + cfg._lock = threading.Lock() + cfg._read() + cfg.consent = True + # Use 100% collection probability so queue-delegation tests are deterministic. + expiry = date.today() + timedelta(days=1) + cfg.set_consent_with_expiry(expiry) + self._cfg = cfg + + worker = _BackgroundWorker(cfg, queue_dir=self._queue_dir) + + self._collector = DataCollector() + self._collector._config = cfg + self._collector._worker = worker + self._collector._init_ok = True + + def tearDown(self) -> None: + shutil.rmtree(str(self._tmp_dir), ignore_errors=True) + + def test_record_delegates_to_worker_queue(self) -> None: + """record() must enqueue the work item without serializing in the caller's thread. + """ + arr = numpy.zeros((256, 256), dtype=numpy.uint16) + q = self._collector._worker._queue + size_before = q.qsize() + self._collector.record("perf_test", "1.0", {"image": arr}) + self.assertEqual( + q.qsize(), + size_before + 1, + "record() must delegate work to the background queue, not serialize inline", + ) + + def test_serialize_creates_zip_with_metadata(self) -> None: + """_serialize() produces a ZIP with valid metadata.json.""" + item = _WorkItem(event_name="zip_test", schema_version="1.0", payload={"score": 0.5}) + self._queue_dir.mkdir(parents=True, exist_ok=True) + zip_path = _serialize(item, self._queue_dir) + self.assertTrue(zip_path.exists()) + with zipfile.ZipFile(str(zip_path)) as zf: + meta = json.loads(zf.read("metadata.json")) + self.assertEqual(meta["event_name"], "zip_test") + self.assertIn("sample_uuid", meta) + self.assertIn("timestamp_utc", meta) + self.assertIn("odemis_version", meta) + + def test_noop_when_consent_false(self) -> None: + """record() does not enqueue work when consent is False.""" + self._cfg.consent = False + q = self._collector._worker._queue + size_before = q.qsize() + self._collector.record("no_consent_test", "1.0", {"x": 1}) + self.assertEqual(q.qsize(), size_before, "No item should be enqueued when consent is False") + + def test_noop_when_consent_none(self) -> None: + """record() does not enqueue work when consent has not been set.""" + self._cfg._cp.remove_option("general", "consent") + q = self._collector._worker._queue + size_before = q.qsize() + self._collector.record("no_consent_none_test", "1.0", {"x": 1}) + self.assertEqual(q.qsize(), size_before, "No item should be enqueued when consent is None") + + def test_no_exception_on_bad_payload_value(self) -> None: + """record() must not raise for unserializable payload values.""" + class _Unserializable: + def __repr__(self): + raise RuntimeError("boom") + + try: + self._collector.record("bad_payload", "1.0", {"bad": _Unserializable()}) + except Exception as exc: + self.fail(f"record() raised unexpectedly: {exc}") + + def test_raises_for_empty_event_name(self) -> None: + """record() raises ValueError for an empty event_name.""" + with self.assertRaises(ValueError): + self._collector.record("", "1.0", {}) + + def test_raises_for_non_string_event_name(self) -> None: + """record() raises ValueError when event_name is not a string.""" + with self.assertRaises(ValueError): + self._collector.record(123, "1.0", {}) # type: ignore[arg-type] + + def test_raises_for_empty_schema_version(self) -> None: + """record() raises ValueError for an empty schema_version.""" + with self.assertRaises(ValueError): + self._collector.record("event", "", {}) + + def test_raises_for_non_dict_payload(self) -> None: + """record() raises ValueError when payload is not a dict.""" + with self.assertRaises(ValueError): + self._collector.record("event", "1.0", [1, 2, 3]) # type: ignore[arg-type] + + def test_raises_for_invalid_image_format(self) -> None: + """record() raises ValueError for an unknown image_format.""" + with self.assertRaises(ValueError): + self._collector.record("event", "1.0", {}, image_format="PNG") + + def test_raises_for_none_image_format(self) -> None: + """record() raises ValueError, not AttributeError, when image_format is None.""" + with self.assertRaises(ValueError): + self._collector.record("event", "1.0", {}, image_format=None) # type: ignore[arg-type] + + def test_raises_for_non_string_image_format(self) -> None: + """record() raises ValueError, not AttributeError, for a non-string image_format.""" + with self.assertRaises(ValueError): + self._collector.record("event", "1.0", {}, image_format=42) # type: ignore[arg-type] + + def test_validation_raises_even_when_consent_false(self) -> None: + """Input validation fires before the consent gate.""" + self._cfg.consent = False + with self.assertRaises(ValueError): + self._collector.record("", "1.0", {}) + + def test_valid_hdf5_format_accepted(self) -> None: + """record() accepts 'HDF5' as image_format without raising.""" + try: + self._collector.record("event", "1.0", {}, image_format="HDF5") + except ValueError as exc: + self.fail(f"record() raised ValueError for valid HDF5 format: {exc}") + +if __name__ == "__main__": + logging.basicConfig(level=logging.DEBUG) + unittest.main() diff --git a/src/odemis/util/test/dc_fetch_test.py b/src/odemis/util/test/dc_fetch_test.py new file mode 100644 index 0000000000..bf9ec66b20 --- /dev/null +++ b/src/odemis/util/test/dc_fetch_test.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +""" +Created on 11 March 2026 + +@author: Karishma Kumar + +Copyright © 2026 Karishma Kumar, Delmic + +This file is part of Odemis. + +Odemis is free software: you can redistribute it and/or modify it under the +terms of the GNU General Public License version 2 as published by the Free +Software Foundation. + +Odemis is distributed in the hope that it will be useful, but WITHOUT ANY +WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR +A PARTICULAR PURPOSE. See the GNU General Public License for more details. + +You should have received a copy of the GNU General Public License along with +Odemis. If not, see http://www.gnu.org/licenses/. +""" + +import tempfile +import unittest +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import Mock, patch + +from odemis.util import dc_fetch + + +class DCFetchTest(unittest.TestCase): + """Unit tests for S3 retrieval helpers.""" + + def test_parse_since_utc_date(self) -> None: + """Date input should parse as UTC midnight.""" + parsed = dc_fetch.parse_since_utc("2026-03-22") + self.assertEqual(parsed, datetime(2026, 3, 22, 0, 0, 0, tzinfo=timezone.utc)) + + def test_parse_key_timestamp(self) -> None: + """Timestamp should be parsed from key basename.""" + parsed = dc_fetch.parse_key_timestamp_utc("host/z_stack_acquired-20260322T104530-a1b2c3d4.zip") + self.assertEqual(parsed, datetime(2026, 3, 22, 10, 45, 30, tzinfo=timezone.utc)) + + def test_should_download_key_filters(self) -> None: + """Event and since filters should both be enforced.""" + key = "host/z_stack_acquired-20260322T104530-a1b2c3d4.zip" + since_before = datetime(2026, 3, 22, 10, 0, 0, tzinfo=timezone.utc) + since_after = datetime(2026, 3, 22, 11, 0, 0, tzinfo=timezone.utc) + self.assertTrue(dc_fetch.should_download_key(key, "z_stack_acquired", since_before)) + self.assertFalse(dc_fetch.should_download_key(key, "other_event", since_before)) + self.assertFalse(dc_fetch.should_download_key(key, "z_stack_acquired", since_after)) + + def test_iter_s3_objects_paginates(self) -> None: + """S3 iterator should follow continuation tokens.""" + client = Mock() + client.list_objects_v2.side_effect = [ + { + "Contents": [{"Key": "host/a.zip"}], + "IsTruncated": True, + "NextContinuationToken": "token-1", + }, + { + "Contents": [{"Key": "host/b.zip"}], + "IsTruncated": False, + }, + ] + keys = [item["Key"] for item in dc_fetch.iter_s3_objects(client, "bucket", "host/")] + self.assertEqual(keys, ["host/a.zip", "host/b.zip"]) + self.assertEqual(client.list_objects_v2.call_count, 2) + + def test_parse_host_filters_comma_list(self) -> None: + """Host parser should accept comma-separated values and normalize them.""" + hosts = dc_fetch.parse_host_filters("meteor-5099, atlas-001 ,/secom-22/") + self.assertEqual(hosts, ["meteor-5099", "atlas-001", "secom-22"]) + + def test_fetch_samples_downloads_matching_keys(self) -> None: + """Fetch flow should download matching keys and report counters.""" + with tempfile.TemporaryDirectory(prefix="dc_fetch_") as tmp_dir: + output_dir = Path(tmp_dir) + client = Mock() + client.list_objects_v2.return_value = { + "Contents": [ + {"Key": "host/evt-20260322T100000-aaaa1111.zip"}, + {"Key": "host/other-20260322T100000-bbbb2222.zip"}, + ], + "IsTruncated": False, + } + + def _download_file(_bucket: str, _key: str, filename: str) -> None: + Path(filename).write_bytes(b"zip") + + client.download_file.side_effect = _download_file + + with patch("odemis.util.dc_fetch.build_s3_client_from_config", return_value=(client, "bucket")): + result = dc_fetch.fetch_samples( + event_filter="evt", + since_utc=datetime(2026, 3, 22, 9, 0, 0, tzinfo=timezone.utc), + output_dir=output_dir, + ) + + self.assertEqual(result["listed"], 2) + self.assertEqual(result["matched"], 1) + self.assertEqual(result["downloaded"], 1) + self.assertEqual(result["failed"], 0) + # Key "host/evt-20260322T100000-aaaa1111.zip" is flattened to + # "host_evt-20260322T100000-aaaa1111.zip" to avoid host collisions. + self.assertTrue((output_dir / "host_evt-20260322T100000-aaaa1111.zip").exists()) + + def test_fetch_samples_applies_host_filter_prefix(self) -> None: + """Host filter should become the S3 list prefix.""" + with tempfile.TemporaryDirectory(prefix="dc_fetch_") as tmp_dir: + output_dir = Path(tmp_dir) + client = Mock() + client.list_objects_v2.return_value = {"Contents": [], "IsTruncated": False} + + with patch("odemis.util.dc_fetch.build_s3_client_from_config", return_value=(client, "bucket")): + dc_fetch.fetch_samples( + event_filter=None, + since_utc=None, + output_dir=output_dir, + host_filter="meteor-5099", + ) + + call_kwargs = client.list_objects_v2.call_args.kwargs + self.assertEqual(call_kwargs["Bucket"], "bucket") + self.assertEqual(call_kwargs["Prefix"], "meteor-5099/") + + def test_fetch_samples_applies_multiple_host_prefixes(self) -> None: + """Comma-separated hosts should trigger one listing call per host prefix.""" + with tempfile.TemporaryDirectory(prefix="dc_fetch_") as tmp_dir: + output_dir = Path(tmp_dir) + client = Mock() + client.list_objects_v2.return_value = {"Contents": [], "IsTruncated": False} + + with patch("odemis.util.dc_fetch.build_s3_client_from_config", return_value=(client, "bucket")): + dc_fetch.fetch_samples( + event_filter=None, + since_utc=None, + output_dir=output_dir, + host_filter="meteor-5099,atlas-001", + ) + + self.assertEqual(client.list_objects_v2.call_count, 2) + first_prefix = client.list_objects_v2.call_args_list[0].kwargs["Prefix"] + second_prefix = client.list_objects_v2.call_args_list[1].kwargs["Prefix"] + self.assertEqual(first_prefix, "meteor-5099/") + self.assertEqual(second_prefix, "atlas-001/") + + def test_fetch_samples_passes_bucket_endpoint_region_overrides(self) -> None: + """Overrides should be forwarded to the S3 client builder.""" + with tempfile.TemporaryDirectory(prefix="dc_fetch_") as tmp_dir: + output_dir = Path(tmp_dir) + client = Mock() + client.list_objects_v2.return_value = {"Contents": [], "IsTruncated": False} + + with patch("odemis.util.dc_fetch.build_s3_client_from_config", return_value=(client, "bucket")) as builder: + dc_fetch.fetch_samples( + event_filter=None, + since_utc=None, + output_dir=output_dir, + host_filter=None, + bucket_override="other-bucket", + endpoint_override="https://s3.eu-west-1.amazonaws.com", + region_override="eu-west-1", + ) + + kwargs = builder.call_args.kwargs + self.assertEqual(kwargs["bucket_override"], "other-bucket") + self.assertEqual(kwargs["endpoint_override"], "https://s3.eu-west-1.amazonaws.com") + self.assertEqual(kwargs["region_override"], "eu-west-1") + + def test_build_s3_client_creates_template_config_when_missing(self) -> None: + """First run should create dc_fetch.ini template and prompt user to fill credentials.""" + with tempfile.TemporaryDirectory(prefix="dc_fetch_cfg_") as tmp_dir: + config_path = Path(tmp_dir) / "dc_fetch.ini" + with self.assertRaises(RuntimeError) as ctx: + dc_fetch.build_s3_client_from_config(config_path=config_path) + + self.assertTrue(config_path.exists()) + self.assertIn(str(config_path), str(ctx.exception)) + content = config_path.read_text(encoding="utf-8") + self.assertIn("[s3]", content) + self.assertIn("access_key =", content) + self.assertIn("secret_key =", content) + + def test_build_s3_client_uses_dc_fetch_ini_values(self) -> None: + """Client builder should read credentials and defaults from dc_fetch.ini.""" + with tempfile.TemporaryDirectory(prefix="dc_fetch_cfg_") as tmp_dir: + config_path = Path(tmp_dir) / "dc_fetch.ini" + config_path.write_text( + "[s3]\n" + "access_key = test-access\n" + "secret_key = test-secret\n" + "bucket = test-bucket\n" + "endpoint_url =\n" + "region = eu-west-1\n", + encoding="utf-8", + ) + + with patch("boto3.client") as mock_boto3_client: + mock_boto3_client.return_value = Mock() + _client, bucket = dc_fetch.build_s3_client_from_config(config_path=config_path) + + self.assertEqual(bucket, "test-bucket") + call_kwargs = mock_boto3_client.call_args.kwargs + self.assertEqual(call_kwargs.get("aws_access_key_id"), "test-access") + self.assertEqual(call_kwargs.get("aws_secret_access_key"), "test-secret") + self.assertEqual(call_kwargs.get("region_name"), "eu-west-1") + self.assertIsNone(call_kwargs.get("endpoint_url")) + + +if __name__ == "__main__": + unittest.main() diff --git a/util/odemis-dc-fetch.py b/util/odemis-dc-fetch.py new file mode 100755 index 0000000000..f5db3ad9ac --- /dev/null +++ b/util/odemis-dc-fetch.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on 11 March 2026 + +@author: Karishma Kumar + +Copyright © 2026 Karishma Kumar, Delmic + +This file is part of Odemis. + +Odemis is free software: you can redistribute it and/or modify it under the +terms of the GNU General Public License version 2 as published by the Free +Software Foundation. + +Odemis is distributed in the hope that it will be useful, but WITHOUT ANY +WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR +A PARTICULAR PURPOSE. See the GNU General Public License for more details. + +You should have received a copy of the GNU General Public License along with +Odemis. If not, see http://www.gnu.org/licenses/. + +Fetch DataCollector ZIP samples from S3. + +Examples: + odemis-dc-fetch + odemis-dc-fetch --output ./downloads + odemis-dc-fetch --event z_stack_acquired + odemis-dc-fetch --since 2026-03-01 + odemis-dc-fetch --host meteor-5099 + odemis-dc-fetch --host meteor-5099,atlas-001,secom-22 + odemis-dc-fetch --bucket delmic-odemis-collect-test --region eu-west-1 + odemis-dc-fetch --since 2026-03-01T12:30:00 --event z_stack_acquired --output ./dc_samples +""" + +import argparse +import logging +import sys +from pathlib import Path +from typing import List, Optional + +from odemis.util.dc_fetch import fetch_samples, parse_since_utc + + +def build_argument_parser() -> argparse.ArgumentParser: + """ + Build CLI argument parser for odemis-dc-fetch. + :return: Configured ArgumentParser instance. + """ + examples = ( + "Examples:\n" + " odemis-dc-fetch\n" + " odemis-dc-fetch --output ./downloads\n" + " odemis-dc-fetch --event z_stack_acquired\n" + " odemis-dc-fetch --since 2026-03-01\n" + " odemis-dc-fetch --host meteor-5099\n" + " odemis-dc-fetch --host meteor-5099,atlas-001,secom-22\n" + " odemis-dc-fetch --bucket delmic-odemis-collect-test --region eu-west-1\n" + " odemis-dc-fetch --since 2026-03-01T12:30:00 --event z_stack_acquired --output ./dc_samples" + ) + parser = argparse.ArgumentParser( + description="Fetch data-collection ZIP samples from S3.", + epilog=examples, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--event", + dest="event", + help="Only fetch samples matching event name.", + ) + parser.add_argument( + "--since", + dest="since", + help="Only fetch samples since UTC date/datetime (e.g. 2026-03-01 or 2026-03-01T12:30:00).", + ) + parser.add_argument( + "--output", + dest="output", + default="./dc_samples", + help="Output directory for downloaded ZIPs (default: ./dc_samples).", + ) + parser.add_argument( + "--host", + dest="host", + help="Optional host/system-id filter; use comma-separated IDs for multiple hosts. " + "By default, fetch across all hosts.", + ) + parser.add_argument( + "--bucket", + dest="bucket", + help="Optional S3 bucket override", + ) + parser.add_argument( + "--endpoint-url", + dest="endpoint_url", + help="Optional S3 endpoint URL override", + ) + parser.add_argument( + "--region", + dest="region", + help="Optional AWS region name for S3 client creation.", + ) + return parser + + +def main(argv: Optional[List[str]] = None) -> int: + """ + Run CLI retrieval flow. + :param argv: Optional command line arguments. + :return: Process return code. + """ + parser = build_argument_parser() + args = parser.parse_args(argv) + + try: + since_utc = parse_since_utc(args.since) if args.since else None + except ValueError: + logging.error("Invalid --since value: %s", args.since) + return 2 + + output_dir = Path(args.output) + try: + result = fetch_samples( + event_filter=args.event, + since_utc=since_utc, + output_dir=output_dir, + host_filter=args.host, + bucket_override=args.bucket, + endpoint_override=args.endpoint_url, + region_override=args.region, + ) + except Exception: + logging.exception("Failed to fetch samples from S3") + return 1 + + print( + "listed={listed} matched={matched} downloaded={downloaded} " + "skipped_existing={skipped_existing} failed={failed}".format(**result) + ) + return 0 if result["failed"] == 0 else 1 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + rc = main(sys.argv[1:]) + logging.shutdown() + sys.exit(rc) diff --git a/util/release-odemis b/util/release-odemis index 93ce690f47..c8944d2a2b 100755 --- a/util/release-odemis +++ b/util/release-odemis @@ -32,6 +32,12 @@ if [ ! -d ~/development/pkg-native/odemis ]; then git clone https://github.com/delmic/odemis.git fi +# Ensure datacollector key is present in the build tree before any release action. +if [ ! -f ~/development/pkg-native/odemis/install/linux/usr/share/odemis/datacollector.key ]; then + echo "Missing required file: ~/development/pkg-native/odemis/install/linux/usr/share/odemis/datacollector.key" + exit 1 +fi + cd ~/development/odemis if ! git remote get-url upstream > /dev/null ; then