dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/registry/project_update_checker.py

@@ -0,0 +1,48 @@
+import logging
+
+from dsgrid.exceptions import DSGInvalidRegistryState
+from .config_update_checker_base import ConfigUpdateCheckerBase
+from .common import DatasetRegistryStatus, ProjectRegistryStatus
+
+
+logger = logging.getLogger(__name__)
+
+
+class ProjectUpdateChecker(ConfigUpdateCheckerBase):
+    """Handles update checks for projects."""
+
+    _ALLOWED_UPDATE_STATUSES = (
+        ProjectRegistryStatus.INITIAL_REGISTRATION,
+        ProjectRegistryStatus.IN_PROGRESS,
+        ProjectRegistryStatus.COMPLETE,
+    )
+    _REQUIRES_DATASET_UNREGISTRATION = (
+        "dimensions",
+        "dimension_mappings",
+    )
+
+    def check_preconditions(self):
+        if self._old_model.status not in self._ALLOWED_UPDATE_STATUSES:
+            msg = f"project status={self._old_model.status} must be one of {self._ALLOWED_UPDATE_STATUSES} in order to update"
+            raise DSGInvalidRegistryState(msg)
+
+    def handle_postconditions(self):
+        # TODO #191: detect changes to required dimensions for each dataset.
+        changes = set(self._REQUIRES_DATASET_UNREGISTRATION).intersection(self._changed_fields)
+        if changes:
+            for dataset in self._new_model.datasets:
+                if dataset.status == DatasetRegistryStatus.REGISTERED:
+                    dataset.status = DatasetRegistryStatus.UNREGISTERED
+            logger.warning(
+                "Set all datasets in %s to unregistered because of changes=%s. "
+                "They must be re-submitted.",
+                self._new_model.project_id,
+                changes,
+            )
+            if self._new_model.status == ProjectRegistryStatus.COMPLETE:
+                self._new_model.status = ProjectRegistryStatus.IN_PROGRESS
+                logger.warning(
+                    "Set project status to %s because of changes=%s.",
+                    self._new_model.status,
+                    changes,
+                )

dsgrid/registry/registration_context.py

@@ -0,0 +1,223 @@
+import getpass
+import logging
+from datetime import datetime
+from typing import Self
+
+from sqlalchemy import Connection
+
+from dsgrid.exceptions import DSGInvalidParameter
+from dsgrid.spark.functions import drop_temp_tables_and_views
+from dsgrid.registry.common import RegistrationModel, RegistryType, VersionUpdateType
+from dsgrid.registry.registry_interface import RegistryInterfaceBase
+from dsgrid.utils.timing import timer_stats_collector, track_timing
+
+
+logger = logging.getLogger(__name__)
+
+
+class RegistrationContext:
+    """Maintains state information across a multi-config registration process."""
+
+    def __init__(
+        self,
+        db: RegistryInterfaceBase,
+        log_message: str,
+        update_type: VersionUpdateType,
+        submitter: str | None,
+    ):
+        self._conn: Connection | None = None
+        self._db = db
+        self._registration = RegistrationModel(
+            timestamp=datetime.now(),
+            submitter=submitter or getpass.getuser(),
+            log_message=log_message,
+            update_type=update_type,
+        )
+        self._managers: dict[RegistryType, RegistryManagerContext | None] = {
+            # This order is required for cleanup in self.finalize().
+            RegistryType.PROJECT: None,
+            RegistryType.DATASET: None,
+            RegistryType.DIMENSION_MAPPING: None,
+            RegistryType.DIMENSION: None,
+        }
+
+    def __del__(self):
+        for registry_type in RegistryType:
+            manager = self._managers.get(registry_type)
+            if manager is not None:
+                logger.warning(
+                    "RegistrationContext destructed with a reference to %s manager",
+                    registry_type.value,
+                )
+                if not manager.offline_mode and manager.has_lock():
+                    logger.error(
+                        "RegistrationContext destructed with a lock on the remote registry. "
+                        "Please contact the dsgrid team. Type=%s IDs=%s",
+                        registry_type.value,
+                        manager.ids,
+                    )
+
+    def __enter__(self) -> Self:
+        self._conn = self._db.engine.connect()
+        self._registration = self._db.insert_registration(self._conn, self._registration)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        if self._conn is None:
+            return
+        try:
+            if exc_type is None:
+                self.finalize(False)
+                self._conn.commit()
+            else:
+                # Order is important. Don't rollback the configs until dataset files are deleted.
+                self.finalize(True)
+                self._conn.rollback()
+        finally:
+            self._conn.close()
+
+    @property
+    def connection(self) -> Connection:
+        """Return the active sqlalchemy connection."""
+        assert self._conn is not None
+        return self._conn
+
+    @property
+    def registration(self) -> RegistrationModel:
+        """Return the registration entry for this context."""
+        return self._registration
+
+    def add_id(self, registry_type: RegistryType, config_id: str, manager):
+        """Add a config ID that has been registered.
+
+        Parameters
+        ----------
+        registry_type : RegistryType
+        config_id : str
+        manager : RegistryManagerBase
+
+        Raises
+        ------
+        DSGInvalidParameter
+            Raised if the config ID is already stored.
+
+        """
+        self.add_ids(registry_type, [config_id], manager)
+
+    def add_ids(self, registry_type: RegistryType, config_ids: list[str], manager):
+        """Add multiple config IDs that have been registered.
+
+        Parameters
+        ----------
+        registry_type : RegistryType
+        config_ids : list[str]
+        manager : RegistryManagerBase
+
+        Raises
+        ------
+        DSGInvalidParameter
+            Raised if a config ID is already stored.
+
+        """
+        manager_context = self._managers[registry_type]
+        if manager_context is None:
+            manager_context = RegistryManagerContext(manager)
+            self._managers[registry_type] = manager_context
+            # manager.acquire_registry_locks(config_ids)
+            # manager_context.set_locked()
+
+        diff = set(config_ids).intersection(manager_context.ids)
+        if diff:
+            msg = f"One or more config IDs are already tracked: {registry_type} {diff}"
+            raise DSGInvalidParameter(msg)
+
+        logger.debug("Added registered IDs: %s %s", registry_type, config_ids)
+        manager_context.ids += config_ids
+
+    def get_ids(self, registry_type: RegistryType):
+        """Return the config IDs for registry_type that have been registered with this context.
+
+        Parameters
+        ----------
+        registry_type : RegistryType
+
+        Returns
+        -------
+        list[str]
+
+        """
+        manager_context = self._managers[registry_type]
+        assert manager_context is not None, registry_type
+        return manager_context.ids
+
+    @track_timing(timer_stats_collector)
+    def finalize(self, error_occurred: bool):
+        """Perform final registration actions. If successful, sync all newly-registered configs
+        and data with the remote registry. If there was an error, remove all intermediate
+        registrations.
+        """
+        try:
+            drop_temp_tables_and_views()
+            for registry_type, manager_context in self._managers.items():
+                if manager_context is not None:
+                    if manager_context.ids:
+                        manager_context.manager.finalize_registration(
+                            self._conn, set(manager_context.ids), error_occurred
+                        )
+                        manager_context.ids.clear()
+                        # manager_context.set_unlocked()
+                    self._managers[registry_type] = None
+        except Exception:
+            logger.exception(
+                "An unexpected error occurred in finalize_registration. "
+                "Please notify the dsgrid team because registry recovery may be required."
+            )
+            raise
+
+
+class RegistryManagerContext:
+    """Maintains state for one registry type."""
+
+    def __init__(self, manager):
+        self._manager = manager
+        self._has_lock = False
+        self._ids = []
+
+    def has_lock(self) -> bool:
+        """Return True if the manager has acquired a lock on the remote registry."""
+        return self._has_lock
+
+    def set_locked(self):
+        """Call when a lock has been acquired on the remote registry."""
+        logger.debug("Locks acquired on remote registry for %s", self._manager.__class__.__name__)
+        self._has_lock = True
+
+    def set_unlocked(self):
+        """Call when all locks have been released on the remote registry."""
+        logger.debug("Locks released on remote registry for %s", self._manager.__class__.__name__)
+        self._has_lock = False
+
+    @property
+    def ids(self):
+        """Return a list of config IDs being managed."""
+        return self._ids
+
+    @ids.setter
+    def ids(self, val):
+        """Return a list of config IDs being managed."""
+        self._ids = val
+
+    @property
+    def manager(self):
+        """Return a RegistryManagerBase"""
+        return self._manager
+
+    @manager.setter
+    def manager(self, val):
+        """Set the RegistryManagerBase"""
+        self._manager = val
+
+    @property
+    def offline_mode(self):
+        """Return True if the manager is in offline mode."""
+        return self._manager.offline_mode

dsgrid/registry/registry_auto_updater.py

@@ -0,0 +1,316 @@
+import logging
+from typing import Any
+
+from dsgrid.config.mapping_tables import MappingTableConfig, MappingTableModel
+from dsgrid.config.dataset_config import DatasetConfig, DatasetConfigModel
+from dsgrid.config.dimension_config import DimensionConfig
+from dsgrid.config.project_config import ProjectConfig
+from dsgrid.exceptions import DSGInvalidParameter
+from dsgrid.registry.registration_context import RegistrationContext
+from dsgrid.registry.common import (
+    ConfigKey,
+    RegistryType,
+    VersionUpdateType,
+)
+from dsgrid.registry.registry_manager import RegistryManager
+
+
+logger = logging.getLogger(__name__)
+
+
+class RegistryAutoUpdater:
+    """Performs auto-updates on the registry."""
+
+    def __init__(self, manager: RegistryManager) -> None:
+        self._project_mgr = manager.project_manager
+        self._dataset_mgr = manager.dataset_manager
+        self._dimension_mgr = manager.dimension_manager
+        self._dimension_mapping_mgr = manager.dimension_mapping_manager
+        self._db = self._dimension_mapping_mgr.db
+
+    def update_dependent_configs(
+        self,
+        config,
+        original_version: str,
+        update_type: VersionUpdateType,
+        log_message: str,
+        submitter: str | None = None,
+    ):
+        """Update all configs that consume this config. Recursive.
+        This is an in incomplete, experimental feature, and is subject to change.
+        Should only be called by an admin that understands the consequences.
+        Passing a dimension may trigger an update to a project and a dimension mapping.
+        The change to that dimension mapping may trigger another update to the project.
+        This guarantees that each config version will only be bumped once.
+
+        It is up to the caller to ensure changes are synced to the remote registry if not in
+        offline mode.
+
+        Datasets likely need to be resubmitted to their projects.
+
+        Parameters
+        ----------
+        config : ConfigBase
+        original_version : str
+            Original version of the config
+        update_type : VersionUpdateType
+        log_message : str
+        """
+        with RegistrationContext(
+            self._project_mgr.db, log_message, update_type, submitter
+        ) as context:
+            return self.update_dependent_configs_with_context(config, original_version, context)
+
+    def update_dependent_configs_with_context(
+        self, config, original_version: str, context: RegistrationContext
+    ):
+        if isinstance(config, DimensionConfig):
+            self._update_dimension_users(config, original_version, context)
+        elif isinstance(config, MappingTableConfig):
+            self._update_dimension_mapping_users(config, original_version, context)
+        elif isinstance(config, DatasetConfig):
+            self._update_dataset_users(config, original_version, context)
+        else:
+            msg = f"Updates of configs dependent on {type(config)}"
+            raise NotImplementedError(msg)
+
+    def _update_dimension_users(
+        self,
+        config: DimensionConfig,
+        original_version: str,
+        context: RegistrationContext,
+    ):
+        # Order is important because
+        # - dimension mappings may have this dimension.
+        # - datasets may have this dimension.
+        # - projects may have this dimension as well as updated mappings and datasets.
+        new_mappings = {}
+        new_datasets = {}
+
+        if config.model.version == original_version:
+            msg = f"current version cannot be the same as the original: {original_version}"
+            raise DSGInvalidParameter(msg)
+
+        affected = self._db.get_containing_models(
+            context.connection, config.model, version=original_version
+        )
+        for mapping in self._update_dimension_mappings_with_dimensions(
+            affected, config, original_version
+        ):
+            key = ConfigKey(mapping.model.mapping_id, mapping.model.version)
+            new_mapping = self._dimension_mapping_mgr.update_with_context(mapping, context)
+            assert key not in new_mappings
+            new_mappings[key] = new_mapping
+            logger.info(
+                "Updated dimension mapping %s to %s as a result of dimension update",
+                new_mapping.model.mapping_id,
+                new_mapping.model.version,
+            )
+        for dataset in self._update_datasets_with_dimensions(affected, config, original_version):
+            key = ConfigKey(dataset.model.dataset_id, dataset.model.version)
+            new_dataset = self._dataset_mgr.update_with_context(dataset, context)
+            assert key not in new_datasets
+            new_datasets[key] = new_dataset
+            logger.info(
+                "Updated dataset %s to %s as a result of dimension update",
+                new_dataset.model.dataset_id,
+                new_dataset.model.version,
+            )
+
+        self._update_projects(
+            context,
+            dimensions={ConfigKey(config.model.dimension_id, original_version): config},
+            dimension_mappings=new_mappings,
+            datasets=new_datasets,
+        )
+
+    def _update_dimension_mapping_users(
+        self,
+        config: MappingTableConfig,
+        original_version: str,
+        context: RegistrationContext,
+    ) -> None:
+        self._update_projects(
+            context,
+            dimension_mappings={ConfigKey(config.model.mapping_id, original_version): config},
+        )
+
+    def _update_dataset_users(
+        self,
+        config: DatasetConfig,
+        original_version: str,
+        context: RegistrationContext,
+    ) -> None:
+        self._update_projects(
+            context,
+            datasets={ConfigKey(config.model.dataset_id, original_version): config},
+        )
+
+    def _update_dimension_mappings_with_dimensions(
+        self,
+        affected: dict[RegistryType, list[Any]],
+        dim: DimensionConfig,
+        original_version: str,
+    ) -> list[MappingTableConfig]:
+        mapping_updates: list[MappingTableConfig] = []
+        for model in affected[RegistryType.DIMENSION_MAPPING]:
+            assert isinstance(model, MappingTableModel)
+            updated = False
+            config = self._dimension_mapping_mgr.get_by_id(model.mapping_id, model.version)
+            if (
+                config.model.from_dimension.dimension_id == dim.model.dimension_id
+                and config.model.from_dimension.version == original_version
+            ):
+                config.model.from_dimension.version = dim.model.version
+                updated = True
+            elif (
+                config.model.to_dimension.dimension_id == dim.model.dimension_id
+                and config.model.from_dimension.version == original_version
+            ):
+                config.model.to_dimension.version = dim.model.version
+                updated = True
+            if updated:
+                mapping_updates.append(config)
+
+        return mapping_updates
+
+    def _update_datasets_with_dimensions(
+        self,
+        affected: dict[RegistryType, list[Any]],
+        dim: DimensionConfig,
+        original_version: str,
+    ) -> list[DatasetConfig]:
+        new_datasets = []
+        for model in affected[RegistryType.DATASET]:
+            assert isinstance(model, DatasetConfigModel)
+            updated = False
+            config = self._dataset_mgr.get_by_id(model.dataset_id, model.version)
+            for ref in config.model.dimension_references:
+                if ref.dimension_id == dim.model.dimension_id and ref.version == original_version:
+                    ref.version = dim.model.version
+                    updated = True
+            if updated:
+                new_datasets.append(config)
+
+        return new_datasets
+
+    def _update_projects(
+        self,
+        context: RegistrationContext,
+        dimensions: dict[ConfigKey, DimensionConfig] | None = None,
+        dimension_mappings: dict[ConfigKey, MappingTableConfig] | None = None,
+        datasets: dict[ConfigKey, DatasetConfig] | None = None,
+    ) -> None:
+        updated_projects = {}
+        if dimensions is not None:
+            self._update_projects_with_new_dimensions(updated_projects, dimensions, context)
+        if dimension_mappings is not None:
+            self._update_projects_with_new_dimension_mappings(
+                updated_projects, dimension_mappings, context
+            )
+        if datasets is not None:
+            self._update_projects_with_new_datasets(updated_projects, datasets, context)
+
+        for project_config in updated_projects.values():
+            new_project = self._project_mgr.update_with_context(project_config, context)
+            logger.info(
+                "Updated project %s to %s as a result of dependent config updates.",
+                new_project.model.project_id,
+                new_project.model.version,
+            )
+        # TODO: Re-submit changed datasets to projects. dataset-to-project mappings might
+        # take some work.
+
+    def _update_projects_with_new_dimensions(
+        self,
+        updated_projects: dict[str, ProjectConfig],
+        dimensions: dict[ConfigKey, DimensionConfig],
+        context: RegistrationContext,
+    ) -> None:
+        """Updates the latest project configurations in place if they consume the dimensions.
+        Edits updated_projects as necessary.
+        """
+        for key, dim in dimensions.items():
+            for model in self._db.get_containing_models(
+                context.connection,
+                dim.model,
+                version=key.version,
+                parent_model_type=RegistryType.PROJECT,
+            )[RegistryType.PROJECT]:
+                config = updated_projects.get(
+                    model.project_id, self._project_mgr.get_by_id(model.project_id)
+                )
+                updated = False
+                for ref in config.model.dimensions.base_dimension_references:
+                    if ref.dimension_id == dim.model.dimension_id and ref.version == key.version:
+                        ref.version = dim.model.version
+                        updated = True
+                        break
+                for ref in config.model.dimensions.supplemental_dimension_references:
+                    if ref.dimension_id == dim.model.dimension_id and ref.version == key.version:
+                        ref.version = dim.model.version
+                        updated = True
+                        break
+                if updated and config.model.project_id not in updated_projects:
+                    updated_projects[config.model.project_id] = config
+
+    def _update_projects_with_new_dimension_mappings(
+        self,
+        updated_projects: dict[str, ProjectConfig],
+        mappings: dict[ConfigKey, MappingTableConfig],
+        context: RegistrationContext,
+    ) -> None:
+        """Updates the latest project configurations in place if they consume the mappings.
+        Edits updated_projects as necessary.
+        """
+        for key, mapping in mappings.items():
+            for model in self._db.get_containing_models(
+                context.connection,
+                mapping.model,
+                version=key.version,
+                parent_model_type=RegistryType.PROJECT,
+            )[RegistryType.PROJECT]:
+                config = updated_projects.get(
+                    model.project_id, self._project_mgr.get_by_id(model.project_id)
+                )
+                updated = False
+                for ref in config.model.dimension_mappings.base_to_supplemental_references:
+                    if ref.mapping_id == mapping.model.mapping_id and ref.version == key.version:
+                        ref.version = mapping.model.version
+                        updated = True
+                        break
+                for ref in config.model.dimensions.supplemental_dimension_references:
+                    if ref.dimension_id == mapping.model.mapping_id and ref.version == key.version:
+                        ref.version = mapping.model.version
+                        updated = True
+                        break
+                if updated and config.model.project_id not in updated_projects:
+                    updated_projects[config.model.project_id] = config
+
+    def _update_projects_with_new_datasets(
+        self,
+        updated_projects: dict[str, ProjectConfig],
+        datasets: dict[ConfigKey, DatasetConfig],
+        context: RegistrationContext,
+    ) -> None:
+        for key, dataset in datasets.items():
+            for model in self._db.get_containing_models(
+                context.connection,
+                dataset.model,
+                version=key.version,
+                parent_model_type=RegistryType.PROJECT,
+            )[RegistryType.PROJECT]:
+                config = updated_projects.get(
+                    model.project_id, self._project_mgr.get_by_id(model.project_id)
+                )
+                updated = False
+                for dataset_ in config.model.datasets:
+                    if (
+                        dataset_.dataset_id == dataset.model.dataset_id
+                        and dataset_.version == key.version
+                    ):
+                        dataset_.version = dataset.model.version
+                        updated = True
+                        break
+                if updated and config.model.project_id not in updated_projects:
+                    updated_projects[config.model.project_id] = config