From 8f1218df6f58570485be04a6740886dc5993f99d Mon Sep 17 00:00:00 2001 From: John Chilton Date: Mon, 8 Apr 2024 13:52:50 -0400 Subject: [PATCH] Work on types for SA 2.0. --- lib/galaxy/app.py | 2 +- lib/galaxy/managers/object_store_instances.py | 16 ++++++----- lib/galaxy/model/__init__.py | 27 +++++++++++++++---- lib/galaxy/model/mapping.py | 6 ++--- 4 files changed, 36 insertions(+), 15 deletions(-) diff --git a/lib/galaxy/app.py b/lib/galaxy/app.py index 304b3ad443c4..c8f3334e9f5c 100644 --- a/lib/galaxy/app.py +++ b/lib/galaxy/app.py @@ -267,7 +267,7 @@ def __init__(self, fsmon=False, **kwargs) -> None: self.vault = self._register_singleton(Vault, VaultFactory.from_app(self)) # type: ignore[type-abstract] self._configure_object_store(fsmon=True) self._register_singleton(BaseObjectStore, self.object_store) - galaxy.model.Dataset.object_store = self.object_store # type:ignore[assignment] + galaxy.model.setup_global_object_store_for_models(self.object_store) def configure_fluent_log(self): if self.config.fluent_log: diff --git a/lib/galaxy/managers/object_store_instances.py b/lib/galaxy/managers/object_store_instances.py index 3031d17f0c1a..bc54560dfd46 100644 --- a/lib/galaxy/managers/object_store_instances.py +++ b/lib/galaxy/managers/object_store_instances.py @@ -26,6 +26,7 @@ from galaxy.model import ( User, UserObjectStore, + OBJECT_STORE_TEMPLATE_CONFIGURATION_VARIABLES_TYPE, ) from galaxy.model.scoped_session import galaxy_scoped_session from galaxy.objectstore import ( @@ -125,19 +126,21 @@ def _upgrade_instance( template = catalog.find_template_by(persisted_object_store.object_store_template_id, payload.template_version) persisted_object_store.object_store_template_version = template.version persisted_object_store.object_store_template_definition = template.model_dump() - old_variables = persisted_object_store.object_store_template_variables + old_variables = persisted_object_store.object_store_template_variables or {} updated_variables = payload.variables - actual_variables = {} + actual_variables: OBJECT_STORE_TEMPLATE_CONFIGURATION_VARIABLES_TYPE = {} for variable in template.variables or []: variable_name = variable.name old_value = old_variables.get(variable_name) - actual_variables[variable_name] = updated_variables.get(variable_name, old_value) + updated_value = updated_variables.get(variable_name, old_value) + if updated_value: + actual_variables[variable_name] = updated_value persisted_object_store.object_store_template_variables = actual_variables - old_secrets = persisted_object_store.object_store_template_secrets + old_secrets = persisted_object_store.object_store_template_secrets or [] new_secrets = payload.secrets - recorded_secrets = persisted_object_store.object_store_template_secrets + recorded_secrets = persisted_object_store.object_store_template_secrets or [] user_vault = trans.user_vault upgraded_template_secrets = [] @@ -299,7 +302,8 @@ def recover_secrets(user_object_store: UserObjectStore, vault: Vault) -> Dict[st # ones recorded as written in the persisted object, the ones # expected in the catalog, or the ones expected in the definition # persisted. - for secret in user_object_store.object_store_template_secrets: + persisted_secret_names = user_object_store.object_store_template_secrets or [] + for secret in persisted_secret_names: vault_key = user_vault_key(user_object_store, secret) secret_value = user_vault.read_secret(vault_key) assert secret_value diff --git a/lib/galaxy/model/__init__.py b/lib/galaxy/model/__init__.py index 4eb25ccef480..088fd781e707 100644 --- a/lib/galaxy/model/__init__.py +++ b/lib/galaxy/model/__init__.py @@ -213,6 +213,11 @@ STR_TO_STR_DICT = Dict[str, str] +OBJECT_STORE_TEMPLATE_CONFIGURATION_VALUE_TYPE = Union[str, bool, int] +OBJECT_STORE_TEMPLATE_CONFIGURATION_VARIABLES_TYPE = Dict[str, OBJECT_STORE_TEMPLATE_CONFIGURATION_VALUE_TYPE] +OBJECT_STORE_TEMPLATE_CONFIGURATION_SECRET_NAMES_TYPE = List[str] +OBJECT_STORE_TEMPLATE_DEFINITION_TYPE = Dict[str, Any] + class TransformAction(TypedDict): action: str @@ -224,6 +229,9 @@ class TransformAction(TypedDict): type_annotation_map={ Optional[STR_TO_STR_DICT]: JSONType, Optional[TRANSFORM_ACTIONS]: MutableJSONType, + Optional[OBJECT_STORE_TEMPLATE_CONFIGURATION_VARIABLES_TYPE]: JSONType, + Optional[OBJECT_STORE_TEMPLATE_CONFIGURATION_SECRET_NAMES_TYPE]: JSONType, + Optional[OBJECT_STORE_TEMPLATE_DEFINITION_TYPE]: JSONType, }, ) @@ -3991,6 +3999,10 @@ def flush(self): sa_session.commit() +def setup_global_object_store_for_models(object_store): + Dataset.object_store = object_store # type:ignore[assignment] + + class Dataset(Base, StorableObject, Serializable): __tablename__ = "dataset" @@ -10973,18 +10985,22 @@ class UserObjectStore(Base, RepresentById): # the id/version and not record the definition... as the templates change # over time this choice has some big consequences despite being easy to swap # implementations. - object_store_template_definition: Mapped[Optional[bytes]] = mapped_column(JSONType) + object_store_template_definition: Mapped[Optional[OBJECT_STORE_TEMPLATE_DEFINITION_TYPE]] = mapped_column(JSONType) # Big JSON blob of the variable name -> value mapping defined for the store's # variables by the user. - object_store_template_variables: Mapped[Optional[bytes]] = mapped_column(JSONType) + object_store_template_variables: Mapped[Optional[OBJECT_STORE_TEMPLATE_CONFIGURATION_VARIABLES_TYPE]] = ( + mapped_column(JSONType) + ) # Track a list of secrets that were defined for this object store at creation - object_store_template_secrets: Mapped[Optional[bytes]] = mapped_column(JSONType) + object_store_template_secrets: Mapped[Optional[OBJECT_STORE_TEMPLATE_CONFIGURATION_SECRET_NAMES_TYPE]] = ( + mapped_column(JSONType) + ) user = relationship("User", back_populates="object_stores") @property def template(self) -> ObjectStoreTemplate: - return ObjectStoreTemplate(**self.object_store_template_definition) + return ObjectStoreTemplate(**self.object_store_template_definition or {}) def object_store_configuration(self, secrets: Dict[str, Any]) -> ObjectStoreConfiguration: user = self.user @@ -10993,9 +11009,10 @@ def object_store_configuration(self, secrets: Dict[str, Any]) -> ObjectStoreConf "email": user.email, "id": user.id, } + variables: OBJECT_STORE_TEMPLATE_CONFIGURATION_VARIABLES_TYPE = self.object_store_template_variables or {} return template_to_configuration( self.template, - variables=self.object_store_template_variables, + variables=variables, secrets=secrets, user_details=user_details, ) diff --git a/lib/galaxy/model/mapping.py b/lib/galaxy/model/mapping.py index ca3e924ba3b9..304a1a18e9a2 100644 --- a/lib/galaxy/model/mapping.py +++ b/lib/galaxy/model/mapping.py @@ -32,7 +32,6 @@ def init( create_tables=False, map_install_models=False, database_query_profiling_proxy=False, - object_store=None, trace_logger=None, use_pbkdf2=True, slow_query_log_threshold=0, @@ -60,7 +59,7 @@ def init( install_mapping.create_database_objects(engine) # Configure model, build ModelMapping - return configure_model_mapping(file_path, object_store, use_pbkdf2, engine, map_install_models, thread_local_log) + return configure_model_mapping(file_path, use_pbkdf2, engine, map_install_models, thread_local_log) def create_additional_database_objects(engine): @@ -105,11 +104,12 @@ def init_models_from_config( config.database_engine_options, map_install_models=map_install_models, database_query_profiling_proxy=config.database_query_profiling_proxy, - object_store=object_store, trace_logger=trace_logger, use_pbkdf2=config.get_bool("use_pbkdf2", True), slow_query_log_threshold=config.slow_query_log_threshold, thread_local_log=config.thread_local_log, log_query_counts=config.database_log_query_counts, ) + if object_store: + model.Dataset.setup_global_object_store_for_models(object_store) return model