dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/query/report_factory.py

@@ -0,0 +1,19 @@
+"""Constructs a report class from a type."""
+
+from dsgrid.query.models import ReportType
+from dsgrid.query.report_peak_load import PeakLoadReport
+from dsgrid.query.reports_base import ReportsBase
+
+
+_TYPE_TO_CLASS = {
+    ReportType.PEAK_LOAD: PeakLoadReport,
+}
+
+
+def make_report(report_type: ReportType) -> ReportsBase:
+    """Make a report class from a report_type."""
+    cls = _TYPE_TO_CLASS.get(report_type)
+    if cls is None:
+        msg = str(report_type)
+        raise NotImplementedError(msg)
+    return cls()

dsgrid/query/report_peak_load.py

@@ -0,0 +1,70 @@
+import logging
+from pathlib import Path
+
+from dsgrid.common import VALUE_COLUMN
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.dataset.models import ValueFormat
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.exceptions import DSGInvalidQuery
+from dsgrid.query.models import ProjectQueryModel
+from dsgrid.spark.functions import join_multiple_columns
+from dsgrid.spark.types import F
+from dsgrid.utils.dataset import ordered_subset_columns
+from dsgrid.utils.files import delete_if_exists
+from dsgrid.utils.spark import read_dataframe
+from .query_context import QueryContext
+from .reports_base import ReportsBase
+
+
+logger = logging.getLogger(__name__)
+
+
+class PeakLoadInputModel(DSGBaseModel):
+    group_by_columns: list[str]
+
+
+class PeakLoadReport(ReportsBase):
+    """Find peak load in a derived dataset."""
+
+    REPORT_FILENAME = "peak_load.parquet"
+
+    def check_query(self, query: ProjectQueryModel) -> None:
+        if query.result.table_format.format_type != ValueFormat.STACKED:
+            msg = "The PeakLoadReport requires the value format to be stacked."
+            raise DSGInvalidQuery(msg)
+
+    def generate(
+        self,
+        filename: Path,
+        output_dir: Path,
+        context: QueryContext,
+        inputs: PeakLoadInputModel,
+    ) -> Path:
+        value_columns = [VALUE_COLUMN]
+        metric_columns = context.get_dimension_column_names(DimensionType.METRIC)
+        if len(metric_columns) > 1:
+            msg = f"Bug: {metric_columns=}"
+            raise Exception(msg)
+        metric_column = next(iter(metric_columns))
+        group_by_columns = inputs.group_by_columns[:]
+        if metric_column not in group_by_columns:
+            group_by_columns.append(metric_column)
+
+        df = read_dataframe(filename)
+        expr = [F.max(x).alias(x) for x in value_columns]
+        peak_load = df.groupBy(*group_by_columns).agg(*expr)
+        join_cols = group_by_columns + value_columns
+        time_columns = context.get_dimension_column_names(DimensionType.TIME)
+        diff = time_columns.difference(df.columns)
+        if diff:
+            msg = f"BUG: expected time column(s) {diff} are not present in table"
+            raise Exception(msg)
+        columns = ordered_subset_columns(df, time_columns) + join_cols
+        with_time = join_multiple_columns(peak_load, df.select(*columns), join_cols).sort(
+            *group_by_columns
+        )
+        output_file = output_dir / PeakLoadReport.REPORT_FILENAME
+        delete_if_exists(output_file)
+        with_time.write.parquet(str(output_file))
+        logger.info("Wrote Peak Load Report to %s", output_file)
+        return output_file

dsgrid/query/reports_base.py

@@ -0,0 +1,20 @@
+import abc
+from pathlib import Path
+from typing import Any
+
+from dsgrid.query.models import ProjectQueryModel
+from dsgrid.query.query_context import QueryContext
+
+
+class ReportsBase(abc.ABC):
+    """Base class for pre-defined reports"""
+
+    @abc.abstractmethod
+    def check_query(self, query: ProjectQueryModel) -> None:
+        """Check compatibility of the user query with the report."""
+
+    @abc.abstractmethod
+    def generate(
+        self, filename: Path, output_dir: Path, context: QueryContext, inputs: Any
+    ) -> Path:
+        """Generate the report on df into output_dir."""

dsgrid/registry/bulk_register.py

@@ -0,0 +1,165 @@
+import getpass
+import logging
+import shutil
+from pathlib import Path
+from uuid import uuid4
+
+from dsgrid.config.registration_models import RegistrationModel, RegistrationJournal
+from dsgrid.dimension.base_models import DatasetDimensionRequirements
+from dsgrid.registry.registry_manager import RegistryManager
+from dsgrid.utils.id_remappings import (
+    map_dimension_ids_to_names,
+    map_dimension_names_to_ids,
+    map_dimension_mapping_names_to_ids,
+    replace_dimension_mapping_names_with_current_ids,
+    replace_dimension_names_with_current_ids,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def bulk_register(
+    registry_manager: RegistryManager,
+    registration_file: Path,
+    data_base_dir: Path | None = None,
+    missing_associations_base_dir: Path | None = None,
+    repo_base_dir: Path | None = None,
+    journal_file: Path | None = None,
+    dataset_dimension_requirements: DatasetDimensionRequirements | None = None,
+):
+    """Bulk register projects, datasets, and their dimensions. If any failure occurs, the code
+    records successfully registered project and dataset IDs to a journal file and prints its
+    filename to the console. Users can pass that filename with the --journal-file option to
+    avoid re-registering those projects and datasets on subsequent attempts.
+
+    The JSON/JSON5 filename must match the data model defined by this documentation:
+
+    https://dsgrid.github.io/dsgrid/reference/data_models/project.html#dsgrid.config.registration_models.RegistrationModel
+    """
+    registration = RegistrationModel.from_file(registration_file)
+    tmp_files = []
+    if journal_file is None:
+        journal_file = Path(f"journal__{uuid4()}.json5")
+        journal = RegistrationJournal()
+    else:
+        journal = RegistrationJournal.from_file(journal_file)
+        registration = registration.filter_by_journal(journal)
+    failure_occurred = False
+    try:
+        return _run_bulk_registration(
+            registry_manager,
+            registration,
+            tmp_files,
+            data_base_dir,
+            missing_associations_base_dir,
+            repo_base_dir,
+            journal,
+            dataset_dimension_requirements,
+        )
+    except Exception:
+        failure_occurred = True
+        raise
+    finally:
+        if failure_occurred and journal.has_entries():
+            journal_file.write_text(journal.model_dump_json(indent=2), encoding="utf-8")
+            logger.info(
+                "Recorded successfully registered projects and datasets to %s. "
+                "Pass this file to the `--journal-file` option of this command to skip those IDs "
+                "on subsequent attempts.",
+                journal_file,
+            )
+        elif journal_file.exists():
+            journal_file.unlink()
+            logger.info("Deleted journal file %s after successful registration.", journal_file)
+        for path in tmp_files:
+            path.unlink()
+
+
+def _run_bulk_registration(
+    mgr: RegistryManager,
+    registration: RegistrationModel,
+    tmp_files: list[Path],
+    data_base_dir: Path | None,
+    missing_associations_base_dir: Path | None,
+    base_repo_dir: Path | None,
+    journal: RegistrationJournal,
+    dataset_dimension_requirements: DatasetDimensionRequirements | None,
+):
+    user = getpass.getuser()
+    project_mgr = mgr.project_manager
+    dataset_mgr = mgr.dataset_manager
+    dim_mgr = mgr.dimension_manager
+    dim_mapping_mgr = mgr.dimension_mapping_manager
+
+    if base_repo_dir is not None:
+        for project in registration.projects:
+            if not project.config_file.is_absolute():
+                project.config_file = base_repo_dir / project.config_file
+        for dataset in registration.datasets:
+            if not dataset.config_file.is_absolute():
+                dataset.config_file = base_repo_dir / dataset.config_file
+        for dataset in registration.dataset_submissions:
+            for field in (
+                "dimension_mapping_file",
+                "dimension_mapping_references_file",
+            ):
+                path = getattr(dataset, field)
+                if path is not None and not path.is_absolute():
+                    setattr(dataset, field, base_repo_dir / path)
+
+    for project in registration.projects:
+        assert project.log_message is not None
+        project_mgr.register(project.config_file, user, project.log_message)
+        journal.add_project(project.project_id)
+
+    for dataset in registration.datasets:
+        config_file = None
+        if dataset.replace_dimension_names_with_ids:
+            mappings = map_dimension_names_to_ids(dim_mgr)
+            orig = dataset.config_file
+            config_file = orig.with_stem(orig.name + "__tmp")
+            shutil.copyfile(orig, config_file)
+            tmp_files.append(config_file)
+            replace_dimension_names_with_current_ids(config_file, mappings)
+        else:
+            config_file = dataset.config_file
+
+        assert dataset.log_message is not None
+        dataset_mgr.register(
+            config_file,
+            user,
+            dataset.log_message,
+            data_base_dir=data_base_dir,
+            missing_associations_base_dir=missing_associations_base_dir,
+            requirements=dataset_dimension_requirements,
+        )
+        journal.add_dataset(dataset.dataset_id)
+
+    for dataset in registration.dataset_submissions:
+        refs_file = None
+        if (
+            dataset.replace_dimension_mapping_names_with_ids
+            and dataset.dimension_mapping_references_file is not None
+        ):
+            dim_id_to_name = map_dimension_ids_to_names(mgr.dimension_manager)
+            mappings = map_dimension_mapping_names_to_ids(dim_mapping_mgr, dim_id_to_name)
+            orig = dataset.dimension_mapping_references_file
+            refs_file = orig.with_stem(orig.name + "__tmp")
+            shutil.copyfile(orig, refs_file)
+            tmp_files.append(refs_file)
+            replace_dimension_mapping_names_with_current_ids(refs_file, mappings)
+        else:
+            refs_file = dataset.dimension_mapping_references_file
+
+        assert dataset.log_message is not None
+        project_mgr.submit_dataset(
+            dataset.project_id,
+            dataset.dataset_id,
+            user,
+            dataset.log_message,
+            dimension_mapping_file=dataset.dimension_mapping_file,
+            dimension_mapping_references_file=refs_file,
+            autogen_reverse_supplemental_mappings=dataset.autogen_reverse_supplemental_mappings,
+        )
+        journal.add_submitted_dataset(dataset.dataset_id, dataset.project_id)

dsgrid/registry/common.py

@@ -0,0 +1,287 @@
+"""Common definitions for registry components"""
+
+import logging
+import re
+from collections import namedtuple
+from datetime import datetime
+from enum import StrEnum
+from pathlib import Path
+
+
+from pydantic import Field
+
+from dsgrid.data_models import DSGBaseModel
+from dsgrid.exceptions import DSGInvalidParameter
+from dsgrid.utils.versioning import make_version
+
+
+REGISTRY_LOG_FILE = "dsgrid_registry.log"
+# Allows letters, numbers, underscores, spaces, dashes
+# Allows letters, numbers, underscores, dashes, spaces
+REGEX_VALID_REGISTRY_NAME = re.compile(r"^[\w -]+$")
+# Allows letters, numbers, underscores, dashes
+REGEX_VALID_REGISTRY_CONFIG_ID_LOOSE = re.compile(r"^[\w/-]+$")
+# Allows letters, numbers, underscores.
+# dataset_id cannot start with a number because of uses in DatasetExpressionHandler
+# It's likely a good rule everywhere else.
+REGEX_VALID_REGISTRY_CONFIG_ID_STRICT = re.compile(r"^[a-zA-Z][\w]+$")
+
+REGISTRY_ID_DELIMITER = "__"
+
+logger = logging.getLogger(__name__)
+
+
+def check_config_id_loose(config_id, tag):
+    # Raises ValueError because this is used in Pydantic models.
+    if not REGEX_VALID_REGISTRY_CONFIG_ID_LOOSE.search(config_id):
+        msg = f"{tag} ID={config_id} is invalid. Restricted to letters, numbers, underscores, and dashes."
+        raise ValueError(msg)
+
+
+def check_config_id_strict(config_id, tag):
+    # Raises ValueError because this is used in Pydantic models.
+    if not REGEX_VALID_REGISTRY_CONFIG_ID_STRICT.search(config_id):
+        msg = (
+            f"{tag} ID={config_id} is invalid. Restricted to letters, numbers, and underscores. "
+            "Cannot start with a number."
+        )
+        raise ValueError(msg)
+
+
+class DatabaseConnection(DSGBaseModel):
+    """Input information to connect to a registry database"""
+
+    url: str
+    # There attributes are commented-out because the registry is currently only
+    # supported in SQLite. If/when we add postgres support, these can be added back.
+    # database: str = "dsgrid"
+    # hostname: str = "localhost"
+    # port: int = 8529
+    # username: str = "root"
+    # password: str = DEFAULT_DB_PASSWORD
+
+    # @classmethod
+    # def from_url(cls, url, **kwargs):
+    # """Create a connection from a URL."""
+    # regex = re.compile(r"http://(.*):(\d+)")
+    # match = regex.search(url)
+    # if match is None:
+    #    raise DSGInvalidParameter(f"Invalid URL format: {url}")
+    # hostname = match.group(1)
+    # port = match.group(2)
+    # return cls(hostname=hostname, port=port, **kwargs)
+
+    def get_filename(self) -> Path:
+        """Return the filename from the URL. Only valid for SQLite databases.
+
+        Raises
+        ------
+        DSGInvalidParameter
+            Raised if the URL does not conform to the SQLite format.
+        """
+        # All call sites will need to be changed if/when we support Postgres.
+        filename = self.try_get_filename()
+        if filename is None:
+            msg = (
+                f"Failed to parse '{self.url}' into a SQLite URL. "
+                "The SQLite file path must be specified in the format 'sqlite:///</path/to/db_file.db>'. "
+            )
+            raise DSGInvalidParameter(msg)
+        return filename
+
+    def try_get_filename(self) -> Path | None:
+        """Return the filename from the URL, if file-based, otherwise None."""
+        regex = re.compile(r"sqlite:\/\/\/(.*)")
+        match = regex.search(self.url)
+        if not match:
+            return None
+        return Path(match.group(1))
+
+
+class DataStoreType(StrEnum):
+    """Specifies the type of data store used for the registry."""
+
+    FILESYSTEM = "filesystem"
+    DUCKDB = "duckdb"
+
+
+class RegistryType(StrEnum):
+    """Registry types"""
+
+    DATASET = "dataset"
+    DIMENSION = "dimension"
+    DIMENSION_MAPPING = "dimension_mapping"
+    PROJECT = "project"
+
+
+MODEL_TYPE_TO_ID_FIELD_MAPPING = {
+    RegistryType.PROJECT: "project_id",
+    RegistryType.DATASET: "dataset_id",
+    RegistryType.DIMENSION: "dimension_id",
+    RegistryType.DIMENSION_MAPPING: "mapping_id",
+}
+
+
+class RegistryTables(StrEnum):
+    """Registry tables"""
+
+    KEY_VALUE = "key_value"
+    CURRENT_VERSIONS = "current_versions"
+    MODELS = "models"
+    REGISTRATIONS = "registrations"
+    CONTAINS = "contains"
+
+
+class DatasetRegistryStatus(StrEnum):
+    """Statuses for a dataset within a project"""
+
+    UNREGISTERED = "Unregistered"
+    REGISTERED = "Registered"
+
+
+class ProjectRegistryStatus(StrEnum):
+    """Statuses for a project within the DSGRID registry"""
+
+    INITIAL_REGISTRATION = "Initial Registration"
+    IN_PROGRESS = "In Progress"
+    COMPLETE = "Complete"
+    PUBLISHED = "Published"
+    DEPRECATED = "Deprecated"
+
+
+class VersionUpdateType(StrEnum):
+    """Types of updates that can be made to projects, datasets, and dimensions"""
+
+    # TODO: we need to find general version update types that can be mapped to
+    #   major, minor and patch.
+    # i.e., replace input_dataset, fix project_config,
+    MAJOR = "major"
+    MINOR = "minor"
+    PATCH = "patch"
+
+
+# These keys are used to store references to project/dataset configs and dimensions
+# in dictionaries.
+ConfigKey = namedtuple("ConfigKey", ["id", "version"])
+
+# Convenience container to be shared among the registry managers.
+# Obviates the need to pass parameters to many constructors.
+RegistryManagerParams = namedtuple(
+    "RegistryManagerParams",
+    [
+        "base_path",
+        "remote_path",
+        "use_remote_data",
+        "fs_interface",
+        "cloud_interface",
+        "offline",
+        "scratch_dir",
+    ],
+)
+
+
+class RegistrationModel(DSGBaseModel):
+    """Registration fields required by the ProjectConfig and DatasetConfig"""
+
+    id: int | None = Field(default=None, description="database ID of the registration")
+    timestamp: datetime = Field(
+        title="timestamp",
+        description="Registration timestamp",
+    )
+    submitter: str = Field(
+        title="submitter",
+        description="Username that submitted the registration",
+    )
+    log_message: str | None = Field(
+        default=None,
+        title="log_message",
+        description="Reason for the update",
+    )
+    update_type: VersionUpdateType = Field(
+        title="update_type",
+        description="Type of update",
+    )
+
+
+def get_version_from_filename(filename):
+    """Return the handle and version from a registry file."""
+    regex = re.compile(r"(?P<handle>\w+)-v(?P<version>[\d\.]+).json5")
+    match = regex.search(filename)
+    assert match, filename
+    return match.groupdict("handle"), make_version(match.groupdict("version"))
+
+
+def make_filename_from_version(handle, version):
+    """Make a filename with the handle and version."""
+    return f"{handle}-v{version}.json5"
+
+
+# def update_version(id_handle, update, registry_path):
+#    """Determine registration or project version for registration.
+#
+#    TODO: Current solution is a quick hack. This needs to be better/formalized.
+#        - Need smarter version updating / checks; use semvar packages
+#        - Set to work with some central version (like S3)
+#        - Currently only updating major version
+#        - NOTE: not currently utilitzing the update_type in
+#                RegistrationModel. Could use this to set
+#                major/minor/patch update decisiosns
+#
+#    Args:
+#        registry_type (RegistryType): type of registry (e.g., Project, Dataset)
+#        id_handle (str): ID handle is either the project_id or dataset_id
+#        update (bool): config registration update setting
+#    """
+#
+#    # TODO: remove when done. project path should be set somewhere else
+#    if not os.path.exists(registry_path):
+#        raise ValueError(f"Path does not exist: {registry_path}")
+#
+#    # if config.update is False, then assume major=1, minor=0, patch=0
+#    if not update:
+#        version = VersionInfo(major=1)
+#        registry_file = Path(registry_path) / make_filename_from_version(id_handle, version)
+#        # Raise error if v1.0.0 registry exists for project_id
+#        if os.path.exists(registry_file):
+#            raise ValueError(
+#                f'{registry_type} registry for "{registry_file}" already '
+#                f"exists. If you want to update the project registration"
+#                f" with a new {registry_type} version, then you will need to"
+#                f" set update=True in {registry_type} config. Alternatively, "
+#                f"if you want to initiate a new dsgrid {registry_type}, you "
+#                "will need to specify a new version handle in the "
+#                f"{registry_type} config."
+#            )
+#    # if update is true...
+#    else:
+#        # list existing project registries
+#        existing_versions = []
+#        for f in os.listdir(registry_path):
+#            handle, version = get_version_from_filename(f)
+#            if handle == id_handle:
+#                existing_versions.append(version)
+#        # check for existing project registries
+#        if not existing_versions:
+#            raise ValueError(
+#                "Registration.update=True, however, no updates can be made "
+#                f"because there are no existing registries for {registry_type}"
+#                f" ID = {id_handle}. Check project_id or set "
+#                f"Registration.update=True in the {registry_type} Config."
+#            )
+#        # find the latest registry version
+#        # NOTE: this is currently based on major verison only
+#        last_version = sorted(existing_versions)[-1]
+#        old_project_version = make_filename_from_version(id_handle, last_version)
+#        old_registry_file = os.path.join(registry_path, old_project_version)
+#
+#        # deprecate old project registry
+#        t = deserialize_registry(old_registry_file)
+#        # DT: Can we use an enum here? Spelling/capitalization mistakes could be costly.
+#        # Deprecated is a project status.
+#        t["status"] = "Deprecated"
+#        # DT: can we use version
+#        t["version"] = last_version.bump_major()
+#        # TODO: deserialize_registry should have returned a Pydantic model
+#        serialize_registry(t, make_filename_from_version(id_handle, t["version"]))
+#
+#    return version

dsgrid/registry/config_update_checker_base.py

@@ -0,0 +1,63 @@
+import abc
+import logging
+
+from dsgrid.exceptions import DSGInvalidOperation
+
+
+logger = logging.getLogger(__name__)
+
+
+class ConfigUpdateCheckerBase(abc.ABC):
+    """Base class for updating all config models"""
+
+    def __init__(self, old_model, new_model):
+        self._old_model = old_model
+        self._new_model = new_model
+        assert type(self._old_model) is type(self._new_model)  # noqa: E721
+        self._type = type(self._old_model)
+        self._changed_fields = set()
+
+    def _check_common(self):
+        for field, attrs in self._type.model_fields.items():
+            old = getattr(self._old_model, field)
+            new = getattr(self._new_model, field)
+            if old != new:
+                extra = attrs.json_schema_extra
+                if extra and not extra.get("updateable", True):
+                    msg = f"{self._type}.{field} cannot be updated"
+                    raise DSGInvalidOperation(msg)
+                self._changed_fields.add(field)
+                logger.info("%s %s changed from %s to %s.", self._type, field, old, new)
+        # FUTURE: We could recurse into each dsgrid pydantic model and check each individual
+        # field. Would also need to handle lists and dicts of models.
+        # This would allow more precise control of changed fields and much better logging.
+
+    @abc.abstractmethod
+    def check_preconditions(self):
+        """Check preconditions for performing an update.
+
+        Raises
+        ------
+        DSGInvalidRegistryState
+            Raised if a precondition is violated.
+
+        """
+
+    @abc.abstractmethod
+    def handle_postconditions(self):
+        """Handle any required postconditions."""
+
+    def run(self):
+        """Run all checks.
+
+        Raises
+        ------
+        DSGInvalidOperation
+            Raised if the user is changing an immutable field.
+        DSGInvalidRegistryState
+            Raised if a precondition is violated.
+
+        """
+        self.check_preconditions()
+        self._check_common()
+        self.handle_postconditions()

dsgrid/registry/data_store_factory.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+from dsgrid.registry.common import DataStoreType
+from dsgrid.registry.data_store_interface import DataStoreInterface
+from dsgrid.registry.duckdb_data_store import DuckDbDataStore
+from dsgrid.registry.filesystem_data_store import FilesystemDataStore
+
+
+def make_data_store(
+    base_path: Path, data_store_type: DataStoreType, initialize: bool = False
+) -> DataStoreInterface:
+    """Factory function to create a data store.
+
+    Parameters
+    ----------
+    base_path : Path
+        The base path for the data store.
+    data_store_type : DataStoreType
+        The type of data store to create.
+    initialize : bool
+        Whether to initialize the data store.
+    """
+    match data_store_type:
+        case DataStoreType.FILESYSTEM:
+            cls = FilesystemDataStore
+        case DataStoreType.DUCKDB:
+            cls = DuckDbDataStore
+        case _:
+            msg = f"Unsupported data store type: {data_store_type}"
+            raise NotImplementedError(msg)
+
+    if initialize:
+        return cls.create(base_path)
+    return cls.load(base_path)