grz-common 0.1.0__py3-none-any.whl

grz_common/__init__.py

@@ -0,0 +1,5 @@
+"""
+Common library for grz applications.
+"""
+
+__version__ = "0.1.0"

grz_common/cli/__init__.py

@@ -0,0 +1,79 @@
+"""
+Common click options for the CLI commands.
+"""
+
+from os import sched_getaffinity
+from pathlib import Path
+
+import click
+import platformdirs
+
+DEFAULT_CONFIG_PATH = Path(platformdirs.user_config_dir("grz-cli")) / "config.yaml"
+
+# Aliases for path types for click options
+# Naming convention: {DIR,FILE}_{Read,Write}_{Exists,Create}
+DIR_R_E = click.Path(
+    exists=True,
+    file_okay=False,
+    dir_okay=True,
+    readable=True,
+    writable=False,
+    resolve_path=True,
+)
+DIR_RW_C = click.Path(
+    exists=False,
+    file_okay=False,
+    dir_okay=True,
+    readable=True,
+    writable=True,
+    resolve_path=True,
+)
+FILE_R_E = click.Path(exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True)
+
+submission_dir = click.option(
+    "--submission-dir",
+    metavar="PATH",
+    type=DIR_R_E,
+    required=True,
+    help="Path to the submission directory containing 'metadata/', 'files/', 'encrypted_files/' and 'logs/' directories",
+)
+
+config_file = click.option(
+    "--config-file",
+    metavar="STRING",
+    type=FILE_R_E,
+    required=False,
+    default=DEFAULT_CONFIG_PATH,
+    help="Path to config file",
+)
+
+threads = click.option(
+    "--threads",
+    default=min(len(sched_getaffinity(0)), 4),
+    type=int,
+    show_default=True,
+    help="Number of threads to use for parallel operations",
+)
+
+submission_id = click.option(
+    "--submission-id",
+    required=True,
+    type=str,
+    metavar="STRING",
+    help="S3 submission ID",
+)
+
+output_dir = click.option(
+    "--output-dir",
+    metavar="PATH",
+    type=DIR_RW_C,
+    required=True,
+    default=None,
+    help="Path to the target submission output directory",
+)
+
+output_json = click.option("--json", "output_json", is_flag=True, help="Output JSON for machine-readability.")
+
+show_details = click.option("--details", "show_details", is_flag=True, help="Show more detailed output.")
+
+force = click.option("--force/--no-force", help="Overwrite files and ignore cached results (dangerous!)")

grz_common/constants.py

@@ -0,0 +1,39 @@
+"""Constants for logging configuration, JSON schema validation, and other settings."""
+
+# This will be overridden by the package that imports it
+PACKAGE_ROOT = "grz_common"
+
+LOGGING_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+LOGGING_DATEFMT = "%Y-%m-%d %I:%M %p"
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": True,
+    "formatters": {
+        "standard": {
+            "format": LOGGING_FORMAT,
+            "datefmt": LOGGING_DATEFMT,
+        },
+    },
+    "handlers": {
+        "default": {
+            "level": "INFO",
+            "formatter": "standard",
+            "class": "logging.StreamHandler",
+        },
+    },
+    "loggers": {
+        "": {  # root logger
+            "handlers": ["default"],
+            "level": "WARNING",
+            "propagate": False,
+        },
+        PACKAGE_ROOT: {"handlers": ["default"], "level": "INFO", "propagate": False},
+        "__main__": {  # if __name__ == '__main__'
+            "handlers": ["default"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+    },
+}
+
+TQDM_SMOOTHING: float = 0.0001

grz_common/logging_setup.py

@@ -0,0 +1,60 @@
+"""
+Module: logging_setup
+
+This module provides functions for setting up logging configuration.
+"""
+
+from __future__ import annotations
+
+import logging
+from os import PathLike
+from pathlib import Path
+
+from .constants import LOGGING_DATEFMT, LOGGING_FORMAT, PACKAGE_ROOT
+
+log = logging.getLogger(__name__)
+
+
+def add_filelogger(
+    file_path: str | PathLike | None = None,
+    level: str = "INFO",
+    logger_name: str = PACKAGE_ROOT,
+) -> None:
+    """
+    Add file logging for the specified package.
+
+    This function configures a file logger to capture log messages
+    for the package specified by logger_name. If no file path
+    is provided, a default log file will be created in the user's
+    home directory.
+
+    :param file_path: Optional; the path to the log file. If None,
+                      a default path will be used.
+    :param level: Optional; the logging level. Default is 'INFO'.
+                  Must be a valid logging level name (e.g., 'DEBUG', 'INFO').
+    :param logger_name: Optional; the name of the logger to add the file handler to.
+                        Default is PACKAGE_ROOT.
+    """
+    logger = logging.getLogger(logger_name)
+
+    if file_path is None:
+        default_log_dir = Path.home() / "logs"
+        default_log_dir.mkdir(parents=True, exist_ok=True)
+        file_path = default_log_dir / f"{logger_name}.log"
+        log.warning("No log file path provided, using default: %s", file_path)
+    else:
+        file_path = Path(file_path)
+
+    try:
+        fh = logging.FileHandler(file_path)
+        fh.setLevel(level.upper())
+        fh.setFormatter(logging.Formatter(LOGGING_FORMAT, LOGGING_DATEFMT))
+        logger.addHandler(fh)
+        log.info(
+            "File logger added for %s at %s with level %s.",
+            logger.name,
+            file_path,
+            level.upper(),
+        )
+    except Exception as e:
+        log.error("Failed to add file logger: %s", e)

grz_common/models/__init__.py

@@ -0,0 +1,3 @@
+"""
+Model classes for grz applications.
+"""

grz_common/models/base.py

@@ -0,0 +1,59 @@
+from os import PathLike
+from pathlib import Path
+from typing import Annotated, Self
+
+from pydantic import AfterValidator, BaseModel, ConfigDict
+from pydantic.types import PathType
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+FilePath = Annotated[Path, AfterValidator(lambda v: v.expanduser()), PathType("file")]
+
+
+class IgnoringBaseModel(BaseModel):
+    model_config = ConfigDict(
+        extra="ignore",
+        validate_assignment=True,
+        use_enum_values=True,
+    )
+
+    def to_yaml(self, fd):
+        """Reads the configuration file and validates it against the schema."""
+        import yaml
+
+        yaml.dump(self.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude_defaults=True))
+
+    @classmethod
+    def from_path(cls, path: str | PathLike) -> Self:
+        """Reads the configuration file and validates it against the schema."""
+        import yaml
+
+        with open(path, encoding="utf-8") as f:
+            config = cls(**yaml.safe_load(f))  # noqa:
+
+        return config
+
+
+class IgnoringBaseSettings(BaseSettings):
+    model_config = SettingsConfigDict(
+        extra="ignore",
+        validate_assignment=True,
+        use_enum_values=True,
+        env_nested_delimiter="__",
+        env_prefix="grz_",
+    )
+
+    def to_yaml(self, fd):
+        """Reads the configuration file and validates it against the schema."""
+        import yaml
+
+        yaml.dump(self.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude_defaults=True), fd)
+
+    @classmethod
+    def from_path(cls, path: str | PathLike) -> Self:
+        """Reads the configuration file and validates it against the schema."""
+        import yaml
+
+        with open(path, encoding="utf-8") as f:
+            config = cls(**yaml.safe_load(f))  # noqa:
+
+        return config

grz_common/models/keys.py

@@ -0,0 +1,45 @@
+from typing import Self
+
+from grz_common.models.base import FilePath, IgnoringBaseSettings
+from pydantic import field_validator, model_validator
+
+
+class KeyModel(IgnoringBaseSettings):
+    grz_public_key: str | None = None
+    """
+    The public key of the recipient (the associated GRZ).
+    """
+
+    grz_public_key_path: FilePath | None = None
+    """
+    Path to the crypt4gh public key of the recipient (the associated GRZ).
+    """
+
+    grz_private_key_path: FilePath | None = None
+    """
+    Path to the crypt4gh private key of the recipient (optional).
+    """
+
+    submitter_private_key_path: FilePath | None = None
+    """
+    Path to the submitter's private key (optional).
+    """
+
+    @field_validator("grz_public_key")
+    @classmethod
+    def check_grz_public_key(cls, v):
+        if v is not None and "BEGIN CRYPT4GH PUBLIC KEY" not in v and "END CRYPT4GH PUBLIC KEY" not in v:
+            raise ValueError("Invalid public key format")
+        return v
+
+    @model_validator(mode="after")
+    def validate_grz_public_key(self) -> Self:
+        if self.grz_public_key is None and self.grz_public_key_path is None:
+            raise ValueError("Either grz_public_key or grz_public_key_path must be set.")
+        if self.grz_public_key is not None and self.grz_public_key_path is not None:
+            raise ValueError("Only one of grz_public_key or grz_public_key_path must be set.")
+        return self
+
+
+class KeyConfigModel(IgnoringBaseSettings):
+    keys: KeyModel

grz_common/models/s3.py

@@ -0,0 +1,79 @@
+from pydantic import AnyHttpUrl, AnyUrl
+
+from .base import IgnoringBaseModel, IgnoringBaseSettings
+
+
+class S3Options(IgnoringBaseModel):
+    endpoint_url: AnyHttpUrl
+    """
+    The URL for the S3 service.
+    """
+
+    bucket: str
+    """
+    The name of the S3 bucket.
+    """
+
+    access_key: str | None = None
+    """
+    The access key for the S3 bucket.
+    If undefined, it is read from the AWS_ACCESS_KEY_ID environment variable.
+    """
+
+    secret: str | None = None
+    """
+    The secret key for the S3 bucket.
+    If undefined, it is read from the AWS_SECRET_ACCESS_KEY environment variable.
+    """
+
+    session_token: str | None = None
+    """
+    The session token for temporary credentials (optional).
+    """
+
+    region_name: str | None = None
+    """
+    The region name for the S3 bucket.
+    """
+
+    api_version: str | None = None
+    """
+    The S3 API version.
+    """
+
+    use_ssl: bool = True
+    """
+    Whether to use SSL for S3 operations.
+    """
+
+    proxy_url: AnyUrl | None = None
+    """
+    The proxy URL for S3 operations (optional).
+    """
+
+    request_checksum_calculation: str | None = None
+    """
+    Whether to calculate checksums for S3 request payloads (optional).
+    Valid values are:
+
+    * ``when_supported`` -- When set, a checksum will be calculated for
+      all request payloads of operations modeled with the ``httpChecksum``
+      trait where ``requestChecksumRequired`` is ``true`` or a
+      ``requestAlgorithmMember`` is modeled.
+
+    * ``when_required`` -- When set, a checksum will only be calculated
+      for request payloads of operations modeled with the ``httpChecksum``
+      trait where ``requestChecksumRequired`` is ``true`` or where a
+      ``requestAlgorithmMember`` is modeled and supplied.
+
+    Defaults to None.
+    """
+
+    multipart_chunksize: int = 256 * 1024**2
+    """
+    The size of the chunks to use for multipart uploads in bytes.
+    """
+
+
+class S3ConfigModel(IgnoringBaseSettings):
+    s3: S3Options

grz_common/progress/__init__.py

@@ -0,0 +1,21 @@
+"""Progress tracking and logging functionality."""
+
+from .progress_logging import FileProgressLogger
+from .states import (
+    DecryptionState,
+    DownloadState,
+    EncryptionState,
+    State,
+    UploadState,
+    ValidationState,
+)
+
+__all__ = [
+    "DecryptionState",
+    "DownloadState",
+    "EncryptionState",
+    "FileProgressLogger",
+    "State",
+    "UploadState",
+    "ValidationState",
+]

grz_common/progress/progress_logging.py

@@ -0,0 +1,193 @@
+"""Module for tracking and logging the state of files over time."""
+
+from __future__ import annotations
+
+import copy
+import json
+import typing
+from collections.abc import Callable
+from os import PathLike
+from pathlib import Path
+
+# pyrefly: ignore
+from grz_pydantic_models.submission.metadata.v1 import File as SubmissionFileMetadata
+
+from ..utils.io import read_multiple_json
+from .states import State
+
+Index = tuple[str, float, int]
+
+
+class FileProgressLogger[T: State]:
+    """
+    A class to log and track the state of files over time. It stores file states in a log file of JSON entries
+    and allows querying the state based on the file path and modification time.
+    """
+
+    _index = {"file_path": str, "modification_time": float, "size": int}
+    # mapping of index -> (metadata, data)
+    _file_states: dict[Index, tuple[SubmissionFileMetadata, T]]
+
+    def __init__(self, log_file_path: str | PathLike):
+        """
+        Initializes the FileProgressLogger instance.
+
+        :param fields: Dictionary of field names and their types, representing the state to track.
+        :param log_file_path: Path to the JSON file where file progress will be logged.
+        """
+        self._file_path = Path(log_file_path)
+        self._file_states = {}
+
+        # Read existing file states from the log file
+        self.read()
+
+    def read(self):
+        """
+        Reads the log file and loads the file states into memory.
+
+        :raises ValueError: If the path exists but is not a file.
+        """
+        if self._file_path.exists():
+            if self._file_path.is_file():
+                with open(self._file_path) as fd:
+                    for row_dict in read_multiple_json(fd):
+                        # Get index and cast them to the correct types
+                        index = typing.cast(
+                            Index,
+                            tuple(self._index[k](row_dict[k]) for k in self._index),
+                        )
+                        # Get metadata
+                        metadata = row_dict["metadata"]
+                        # Get state
+                        state = row_dict["state"]
+
+                        self._file_states[index] = (metadata, state)
+            else:
+                raise ValueError(f"Path is not a file: '{str(self._file_path)}'")
+
+    def cleanup(self, keep: list[tuple[PathLike, SubmissionFileMetadata]]):
+        """
+        Removes all entries from the log file and in-memory state that are not in the keep list.
+
+        :param keep: List of tuples containing the file path and metadata of files to keep.
+        """
+        self._file_path.unlink(missing_ok=True)
+        for file, file_metadata in keep:
+            state = self.get_state(file, file_metadata)
+            if state is not None:
+                self.set_state(file, file_metadata, state)
+
+    def _get_index(self, file_path: str | PathLike) -> Index:
+        """
+        Generates a unique index for a given file based on its name and modification time.
+
+        :param file_path: Path object representing the file.
+        :return: A tuple containing the file name and modification time.
+        """
+        file_path = Path(file_path).resolve()
+
+        if file_path.is_file():
+            return str(file_path), file_path.stat().st_mtime, file_path.stat().st_size
+        else:
+            return str(file_path), -1, -1  # catches files that do not exist
+
+    # def get_index(self, file_path: Path, file_metadata: Dict) -> tuple:
+    #     return self._get_index(file_path, file_metadata)
+
+    def get_state(
+        self,
+        file_path: str | PathLike,
+        file_metadata: dict | SubmissionFileMetadata,
+        default: T | Callable[[Path, SubmissionFileMetadata], T] | None = None,
+    ) -> T | None:
+        """
+        Retrieves the stored state of a file if it exists in the log.
+
+        :param file_path: The file path to query for the state.
+        :param file_metadata: The metadata of the file to query for the state.
+        :param default: Default state to use if the file does not exist.
+            Can be a Callable that takes the file path and the file metadata as input and returns some state:
+            `Callable[[Path, SubmissionFileMetadata], T]`.
+
+            The default state gets automatically saved as the state for this file in case there is no stored state.
+        :return: A dictionary representing the file's state, or None if the file's state isn't logged.
+        """
+        file_path = Path(file_path)
+        index = self._get_index(file_path)
+
+        # get stored state
+        stored_metadata, stored_data = self._file_states.get(index, (None, None))
+
+        # check if metadata matches
+        if not isinstance(file_metadata, SubmissionFileMetadata):
+            file_metadata = SubmissionFileMetadata(**file_metadata)
+
+        if stored_metadata and not isinstance(stored_metadata, SubmissionFileMetadata):
+            stored_metadata = SubmissionFileMetadata(**stored_metadata)
+
+        if stored_metadata and file_metadata == stored_metadata:
+            return copy.deepcopy(stored_data)
+
+        # metadata mismatch -> no valid stored state -
+        if file_metadata and default:
+            if callable(default):
+                state = default(file_path, file_metadata)
+                self.set_state(file_path, file_metadata, state)
+                return state
+            else:
+                if not default:
+                    raise ValueError("Default state must be provided if not callable")
+                    # return None
+                default = typing.cast(T, default)
+                self.set_state(file_path, file_metadata, default)
+                return default
+        else:
+            return None
+
+    def set_state(
+        self,
+        file_path: str | PathLike,
+        file_metadata: dict | SubmissionFileMetadata,
+        state: T,
+    ):
+        """
+        Log the state of a file:
+         - Update the in-memory state
+         - Persist the state to the JSON log file
+
+        :param file_path: The path of the file whose state is being set.
+        :param file_metadata: Submission file metadata to store
+        :param state: A dictionary containing the file's state data to be logged.
+        """
+        file_path = Path(file_path)
+        index = self._get_index(file_path)
+
+        if file_metadata and not isinstance(file_metadata, SubmissionFileMetadata):
+            file_metadata = SubmissionFileMetadata(**file_metadata)
+        file_metadata = typing.cast(SubmissionFileMetadata, file_metadata)
+
+        # Update state in memory
+        self._file_states[index] = (file_metadata, state)
+
+        # Persist state to JSON log file
+        with open(self._file_path, "a", newline="") as fd:
+            # Append the new state row to the log file
+            json.dump(
+                {
+                    # index keys
+                    **{k: v for k, v in zip(self._index.keys(), index, strict=True)},
+                    # state
+                    "metadata": file_metadata.model_dump(by_alias=True),
+                    "state": state,
+                },
+                fd,
+            )
+            fd.write("\n")
+
+    def num_entries(self) -> int:
+        """
+        Returns the number of entries in the file_states dictionary
+
+        :return: An integer representing the number of entries in the file_states dictionary
+        """
+        return len(self._file_states)

grz_common/progress/states.py

@@ -0,0 +1,53 @@
+"""
+This module contains the type definitions for the progress logging states.
+"""
+
+from typing import TypedDict
+
+
+class State(TypedDict, total=False):
+    """
+    Base state class for progress logging.
+    """
+
+    errors: list[str]
+
+
+class ValidationState(State):
+    """
+    State for validation progress.
+    """
+
+    validation_passed: bool
+
+
+class UploadState(State):
+    """
+    State for upload progress.
+    """
+
+    upload_successful: bool
+
+
+class EncryptionState(State):
+    """
+    State for encryption progress.
+    """
+
+    encryption_successful: bool
+
+
+class DecryptionState(State):
+    """
+    State for decryption progress.
+    """
+
+    decryption_successful: bool
+
+
+class DownloadState(State):
+    """
+    State for download progress.
+    """
+
+    download_successful: bool