Merge pull request #12638 from jmchilton/better_model_stores

Better model store abstractions.

lib/galaxy/model/__init__.py

@@ -4,6 +4,7 @@ Galaxy data model classes
 Naming: try to use class names that have a distinct plural form so that
 the relationship cardinalities are obvious (e.g. prefer Dataset to Data)
 """
+import abc
 import base64
 import errno
 import json
@@ -20,6 +21,8 @@ from datetime import datetime, timedelta
 from enum import Enum
 from string import Template
 from typing import (
+    Any,
+    Dict,
     Iterable,
     List,
     NamedTuple,
@@ -100,6 +103,7 @@ from galaxy.model.item_attrs import get_item_annotation_str, UsesAnnotations
 from galaxy.model.orm.now import now
 from galaxy.model.view import HistoryDatasetCollectionJobStateSummary
 from galaxy.security import get_permitted_actions
+from galaxy.security.idencoding import IdEncodingHelper
 from galaxy.security.validate_user_input import validate_password_str
 from galaxy.util import (
     directory_hash_id,
@@ -280,6 +284,23 @@ class SerializationOptions:
             self.serialize_files_handler.serialize_files(dataset, as_dict)
 
 
+class Serializable(RepresentById):
+
+    def serialize(self, id_encoder: IdEncodingHelper, serialization_options: SerializationOptions, for_link: bool = False) -> Dict[str, Any]:
+        """Serialize model for a re-population in (potentially) another Galaxy instance."""
+        if for_link:
+            rval = dict_for(
+                self
+            )
+            serialization_options.attach_identifier(id_encoder, self, rval)
+            return rval
+        return self._serialize(id_encoder, serialization_options)
+
+    @abc.abstractmethod
+    def _serialize(self, id_encoder: IdEncodingHelper, serialization_options: SerializationOptions) -> Dict[str, Any]:
+        """Serialize model for a re-population in (potentially) another Galaxy instance."""
+
+
 class HasName:
 
     def get_display_name(self):
@@ -854,7 +875,7 @@ class TaskMetricNumeric(BaseJobMetric, RepresentById):
     metric_value = Column(Numeric(JOB_METRIC_PRECISION, JOB_METRIC_SCALE))
 
 
-class Job(Base, JobLike, UsesCreateAndUpdateTime, Dictifiable, RepresentById):
+class Job(Base, JobLike, UsesCreateAndUpdateTime, Dictifiable, Serializable):
     """
     A job represents a request to run a tool given input datasets, tool
     parameters, and output datasets.
@@ -1306,7 +1327,7 @@ class Job(Base, JobLike, UsesCreateAndUpdateTime, Dictifiable, RepresentById):
             if flush:
                 object_session(self).flush()
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         job_attrs = dict_for(self)
         serialization_options.attach_identifier(id_encoder, self, job_attrs)
         job_attrs['tool_id'] = self.tool_id
@@ -1917,7 +1938,7 @@ class ImplicitlyCreatedDatasetCollectionInput(Base, RepresentById):
         self.input_dataset_collection = input_dataset_collection
 
 
-class ImplicitCollectionJobs(Base, RepresentById):
+class ImplicitCollectionJobs(Base, Serializable):
     __tablename__ = 'implicit_collection_jobs'
 
     id = Column(Integer, primary_key=True)
@@ -1937,7 +1958,7 @@ class ImplicitCollectionJobs(Base, RepresentById):
     def job_list(self):
         return [icjja.job for icjja in self.jobs]
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             populated_state=self.populated_state,
@@ -2288,7 +2309,7 @@ class HistoryAudit(Base, RepresentById):
         sa_session.execute(q)
 
 
-class History(Base, HasTags, Dictifiable, UsesAnnotations, HasName, RepresentById):
+class History(Base, HasTags, Dictifiable, UsesAnnotations, HasName, Serializable):
     __tablename__ = 'history'
     __table_args__ = (
         Index('ix_history_slug', 'slug', mysql_length=200),
@@ -2568,7 +2589,7 @@ class History(Base, HasTags, Dictifiable, UsesAnnotations, HasName, RepresentByI
         # This needs to be a list
         return [hda for hda in self.datasets if not hda.dataset.deleted]
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
 
         history_attrs = dict_for(
             self,
@@ -3094,7 +3115,7 @@ class StorableObject:
             sa_session.flush()
 
 
-class Dataset(StorableObject, RepresentById, _HasTable):
+class Dataset(StorableObject, Serializable, _HasTable):
 
     class states(str, Enum):
         NEW = 'new'
@@ -3340,7 +3361,7 @@ class Dataset(StorableObject, RepresentById, _HasTable):
                 return True
         return False
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         # serialize Dataset objects only for jobs that can actually modify these models.
         assert serialization_options.serialize_dataset_objects
 
@@ -3365,7 +3386,7 @@ class Dataset(StorableObject, RepresentById, _HasTable):
         return rval
 
 
-class DatasetSource(Base, RepresentById):
+class DatasetSource(Base, Serializable):
     __tablename__ = 'dataset_source'
 
     id = Column(Integer, primary_key=True)
@@ -3376,7 +3397,7 @@ class DatasetSource(Base, RepresentById):
     dataset = relationship('Dataset', back_populates='sources')
     hashes = relationship('DatasetSourceHash', back_populates='source')
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             source_uri=self.source_uri,
@@ -3388,7 +3409,7 @@ class DatasetSource(Base, RepresentById):
         return rval
 
 
-class DatasetSourceHash(Base, RepresentById):
+class DatasetSourceHash(Base, Serializable):
     __tablename__ = 'dataset_source_hash'
 
     id = Column(Integer, primary_key=True)
@@ -3397,7 +3418,7 @@ class DatasetSourceHash(Base, RepresentById):
     hash_value = Column(TEXT)
     source = relationship('DatasetSource', back_populates='hashes')
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             hash_function=self.hash_function,
@@ -3407,7 +3428,7 @@ class DatasetSourceHash(Base, RepresentById):
         return rval
 
 
-class DatasetHash(Base, RepresentById):
+class DatasetHash(Base, Serializable):
     __tablename__ = 'dataset_hash'
 
     id = Column(Integer, primary_key=True)
@@ -3417,7 +3438,7 @@ class DatasetHash(Base, RepresentById):
     extra_files_path = Column(TEXT)
     dataset = relationship('Dataset', back_populates='hashes')
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             hash_function=self.hash_function,
@@ -3936,14 +3957,7 @@ class DatasetInstance:
 
         return msg
 
-    def serialize(self, id_encoder, serialization_options, for_link=False):
-        if for_link:
-            rval = dict_for(
-                self
-            )
-            serialization_options.attach_identifier(id_encoder, self, rval)
-            return rval
-
+    def _serialize(self, id_encoder, serialization_options):
         metadata = _prepare_metadata_for_serialization(id_encoder, serialization_options, self.metadata)
         rval = dict_for(
             self,
@@ -3984,7 +3998,7 @@ class DatasetInstance:
 
 
 class HistoryDatasetAssociation(DatasetInstance, HasTags, Dictifiable, UsesAnnotations,
-                                HasName, RepresentById):
+                                HasName, Serializable):
     """
     Resource class that creates a relation between a dataset and a user history.
     """
@@ -4207,15 +4221,8 @@ class HistoryDatasetAssociation(DatasetInstance, HasTags, Dictifiable, UsesAnnot
                 rval += self.get_total_size()
         return rval
 
-    def serialize(self, id_encoder, serialization_options, for_link=False):
-        if for_link:
-            rval = dict_for(
-                self
-            )
-            serialization_options.attach_identifier(id_encoder, self, rval)
-            return rval
-
-        rval = super().serialize(id_encoder, serialization_options)
+    def _serialize(self, id_encoder, serialization_options):
+        rval = super()._serialize(id_encoder, serialization_options)
         rval['state'] = self.state
         rval["hid"] = self.hid
         rval["annotation"] = unicodify(getattr(self, 'annotation', ''))
@@ -4321,7 +4328,7 @@ class HistoryDatasetAssociation(DatasetInstance, HasTags, Dictifiable, UsesAnnot
                  + type_coerce(cls.id, Unicode)).label('type_id'))
 
 
-class HistoryDatasetAssociationHistory(Base, RepresentById):
+class HistoryDatasetAssociationHistory(Base, Serializable):
     __tablename__ = 'history_dataset_association_history'
 
     id = Column(Integer, primary_key=True)
@@ -4397,7 +4404,7 @@ class HistoryDatasetAssociationSubset(Base, RepresentById):
         self.location = location
 
 
-class Library(Base, Dictifiable, HasName, RepresentById):
+class Library(Base, Dictifiable, HasName, Serializable):
     __tablename__ = 'library'
 
     id = Column(Integer, primary_key=True)
@@ -4422,7 +4429,7 @@ class Library(Base, Dictifiable, HasName, RepresentById):
         self.synopsis = synopsis
         self.root_folder = root_folder
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             name=self.name,
@@ -4475,7 +4482,7 @@ class Library(Base, Dictifiable, HasName, RepresentById):
         return roles
 
 
-class LibraryFolder(Base, Dictifiable, HasName, RepresentById):
+class LibraryFolder(Base, Dictifiable, HasName, Serializable):
     __tablename__ = 'library_folder'
     __table_args__ = (
         Index('ix_library_folder_name', 'name', mysql_length=200),
@@ -4552,7 +4559,7 @@ class LibraryFolder(Base, Dictifiable, HasName, RepresentById):
         # This needs to be a list
         return [ld for ld in self.datasets if ld.library_dataset_dataset_association and not ld.library_dataset_dataset_association.dataset.deleted]
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             name=self.name,
@@ -4597,7 +4604,7 @@ class LibraryFolder(Base, Dictifiable, HasName, RepresentById):
         return f.library_root[0]
 
 
-class LibraryDataset(Base, RepresentById):
+class LibraryDataset(Base, Serializable):
     __tablename__ = 'library_dataset'
 
     id = Column(Integer, primary_key=True)
@@ -4664,7 +4671,7 @@ class LibraryDataset(Base, RepresentById):
     def display_name(self):
         self.library_dataset_dataset_association.display_name()
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             name=self.name,
@@ -4714,7 +4721,7 @@ class LibraryDataset(Base, RepresentById):
         return rval
 
 
-class LibraryDatasetDatasetAssociation(DatasetInstance, HasName, RepresentById):
+class LibraryDatasetDatasetAssociation(DatasetInstance, HasName, Serializable):
 
     def __init__(self,
                  copied_from_history_dataset_association=None,
@@ -4805,15 +4812,8 @@ class LibraryDatasetDatasetAssociation(DatasetInstance, HasName, RepresentById):
     def has_manage_permissions_roles(self, security_agent):
         return self.dataset.has_manage_permissions_roles(security_agent)
 
-    def serialize(self, id_encoder, serialization_options, for_link=False):
-        if for_link:
-            rval = dict_for(
-                self
-            )
-            serialization_options.attach_identifier(id_encoder, self, rval)
-            return rval
-
-        rval = super().serialize(id_encoder, serialization_options)
+    def _serialize(self, id_encoder, serialization_options):
+        rval = super()._serialize(id_encoder, serialization_options)
         self._handle_serialize_files(id_encoder, serialization_options, rval)
         return rval
 
@@ -5082,7 +5082,7 @@ class InnerCollectionFilter(NamedTuple):
         return self.operator_function(getattr(table, self.column), self.expected_value)
 
 
-class DatasetCollection(Base, Dictifiable, UsesAnnotations, RepresentById):
+class DatasetCollection(Base, Dictifiable, UsesAnnotations, Serializable):
     __tablename__ = 'dataset_collection'
 
     id = Column(Integer, primary_key=True)
@@ -5393,7 +5393,7 @@ class DatasetCollection(Base, Dictifiable, UsesAnnotations, RepresentById):
     def has_subcollections(self):
         return ":" in self.collection_type
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             type=self.collection_type,
@@ -5462,7 +5462,7 @@ class HistoryDatasetCollectionAssociation(
     HasTags,
     Dictifiable,
     UsesAnnotations,
-    RepresentById
+    Serializable,
 ):
     """ Associates a DatasetCollection with a History. """
     __tablename__ = 'history_dataset_collection_association'
@@ -5597,14 +5597,7 @@ class HistoryDatasetCollectionAssociation(
         if len(rval) > 0:
             return rval if multiple else rval[0]
 
-    def serialize(self, id_encoder, serialization_options, for_link=False):
-        if for_link:
-            rval = dict_for(
-                self
-            )
-            serialization_options.attach_identifier(id_encoder, self, rval)
-            return rval
-
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             display_name=self.display_name(),
@@ -5783,7 +5776,7 @@ class LibraryDatasetCollectionAssociation(Base, DatasetCollectionInstance, Repre
         return dict_value
 
 
-class DatasetCollectionElement(Base, Dictifiable, RepresentById):
+class DatasetCollectionElement(Base, Dictifiable, Serializable):
     """ Associates a DatasetInstance (hda or ldda) with a DatasetCollection. """
     __tablename__ = 'dataset_collection_element'
 
@@ -5927,7 +5920,7 @@ class DatasetCollectionElement(Base, Dictifiable, RepresentById):
         )
         return new_element
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         rval = dict_for(
             self,
             element_type=self.element_type,
@@ -7392,7 +7385,7 @@ class WorkflowInvocationStepOutputDatasetCollectionAssociation(Base, Dictifiable
     dict_collection_visible_keys = ['id', 'workflow_invocation_step_id', 'dataset_collection_id', 'output_name']
 
 
-class MetadataFile(Base, StorableObject, RepresentById):
+class MetadataFile(Base, StorableObject, Serializable):
     __tablename__ = 'metadata_file'
 
     id = Column(Integer, primary_key=True)
@@ -7450,7 +7443,7 @@ class MetadataFile(Base, StorableObject, RepresentById):
             # Return filename inside hashed directory
             return os.path.abspath(os.path.join(path, "metadata_%d.dat" % self.id))
 
-    def serialize(self, id_encoder, serialization_options):
+    def _serialize(self, id_encoder, serialization_options):
         as_dict = dict_for(self)
         serialization_options.attach_identifier(id_encoder, self, as_dict)
         as_dict["uuid"] = str(self.uuid or '') or None

lib/galaxy/model/store/__init__.py

@@ -10,7 +10,7 @@ from json import (
     dumps,
     load,
 )
-from typing import Any, Dict, List
+from typing import Any, cast, Dict, List, Optional, Union
 from uuid import uuid4
 
 from bdbag import bdbag_api as bdb
@@ -28,6 +28,9 @@ from ..custom_types import json_encoder
 from ..item_attrs import add_item_annotation, get_item_annotation_str
 from ... import model
 
+
+ObjectKeyType = Union[str, int]
+
 ATTRS_FILENAME_HISTORY = 'history_attrs.txt'
 ATTRS_FILENAME_DATASETS = 'datasets_attrs.txt'
 ATTRS_FILENAME_JOBS = 'jobs_attrs.txt'
@@ -120,6 +123,11 @@ class ModelImportStore(metaclass=abc.ABCMeta):
         and a history may contain multiple objects with the same 'hid'.
         """
 
+    @property
+    def file_source_root(self) -> Optional[str]:
+        """Source of valid file data."""
+        return None
+
     def trust_hid(self, obj_attrs) -> bool:
         """Trust HID when importing objects into a new History."""
 
@@ -272,6 +280,7 @@ class ModelImportStore(metaclass=abc.ABCMeta):
                 metadata = dataset_attrs['metadata']
 
                 model_class = dataset_attrs.get("model_class", "HistoryDatasetAssociation")
+                dataset_instance: model.DatasetInstance
                 if model_class == "HistoryDatasetAssociation":
                     # Create dataset and HDA.
                     dataset_instance = model.HistoryDatasetAssociation(name=dataset_attrs['name'],
@@ -318,15 +327,17 @@ class ModelImportStore(metaclass=abc.ABCMeta):
                 self._flush()
 
                 if model_class == "HistoryDatasetAssociation":
+                    hda = cast(model.HistoryDatasetAssociation, dataset_instance)
                     # don't use add_history to manage HID handling across full import to try to preserve
                     # HID structure.
-                    dataset_instance.history = history
+                    hda.history = history
                     if new_history and self.trust_hid(dataset_attrs):
-                        dataset_instance.hid = dataset_attrs['hid']
+                        hda.hid = dataset_attrs['hid']
                     else:
-                        object_import_tracker.requires_hid.append(dataset_instance)
+                        object_import_tracker.requires_hid.append(hda)
 
                 self._flush()
+                file_source_root = self.file_source_root
 
                 # If dataset is in the dictionary - we will assert this dataset is tied to the Galaxy instance
                 # and the import options are configured for allowing editing the dataset (e.g. for metadata setting).
@@ -337,15 +348,16 @@ class ModelImportStore(metaclass=abc.ABCMeta):
                 else:
                     file_name = dataset_attrs.get('file_name')
                     if file_name:
+                        assert file_source_root
                         # Do security check and move/copy dataset data.
-                        archive_path = os.path.abspath(os.path.join(self.archive_dir, file_name))
+                        archive_path = os.path.abspath(os.path.join(file_source_root, file_name))
                         if os.path.islink(archive_path):
                             raise MalformedContents(f"Invalid dataset path: {archive_path}")
 
                         temp_dataset_file_name = \
                             os.path.realpath(archive_path)
 
-                        if not in_directory(temp_dataset_file_name, self.archive_dir):
+                        if not in_directory(temp_dataset_file_name, file_source_root):
                             raise MalformedContents(f"Invalid dataset path: {temp_dataset_file_name}")
 
                     if not file_name or not os.path.exists(temp_dataset_file_name):
@@ -361,14 +373,15 @@ class ModelImportStore(metaclass=abc.ABCMeta):
                         # Import additional files if present. Histories exported previously might not have this attribute set.
                         dataset_extra_files_path = dataset_attrs.get('extra_files_path', None)
                         if dataset_extra_files_path:
+                            assert file_source_root
                             dir_name = dataset_instance.dataset.extra_files_path_name
-                            dataset_extra_files_path = os.path.join(self.archive_dir, dataset_extra_files_path)
+                            dataset_extra_files_path = os.path.join(file_source_root, dataset_extra_files_path)
                             for root, _dirs, files in safe_walk(dataset_extra_files_path):
                                 extra_dir = os.path.join(dir_name, root.replace(dataset_extra_files_path, '', 1).lstrip(os.path.sep))
                                 extra_dir = os.path.normpath(extra_dir)
                                 for extra_file in files:
                                     source = os.path.join(root, extra_file)
-                                    if not in_directory(source, self.archive_dir):
+                                    if not in_directory(source, file_source_root):
                                         raise MalformedContents(f"Invalid dataset path: {source}")
                                     self.object_store.update_from_file(
                                         dataset_instance.dataset, extra_dir=extra_dir,
@@ -653,27 +666,9 @@ class ModelImportStore(metaclass=abc.ABCMeta):
     def _import_jobs(self, object_import_tracker, history):
         object_key = self.object_key
 
-        def _find_hda(input_key):
-            hda = None
-            if input_key in object_import_tracker.hdas_by_key:
-                hda = object_import_tracker.hdas_by_key[input_key]
-            if input_key in object_import_tracker.hda_copied_from_sinks:
-                hda = object_import_tracker.hdas_by_key[object_import_tracker.hda_copied_from_sinks[input_key]]
-            return hda
-
-        def _find_hdca(input_key):
-            hdca = None
-            if input_key in object_import_tracker.hdcas_by_key:
-                hdca = object_import_tracker.hdcas_by_key[input_key]
-            if input_key in object_import_tracker.hdca_copied_from_sinks:
-                hdca = object_import_tracker.hdcas_by_key[object_import_tracker.hdca_copied_from_sinks[input_key]]
-            return hdca
-
-        def _find_dce(input_key):
-            dce = None
-            if input_key in object_import_tracker.dces_by_key:
-                dce = object_import_tracker.dces_by_key[input_key]
-            return dce
+        _find_hda = object_import_tracker.find_hda
+        _find_hdca = object_import_tracker.find_hdca
+        _find_dce = object_import_tracker.find_dce
 
         #
         # Create jobs.
@@ -788,6 +783,18 @@ class ObjectImportTracker:
 
     Needed to re-establish connections and such in multiple passes.
     """
+    libraries_by_key: Dict[ObjectKeyType, model.Library]
+    hdas_by_key: Dict[ObjectKeyType, model.HistoryDatasetAssociation]
+    hdas_by_id: Dict[int, model.HistoryDatasetAssociation]
+    hdcas_by_key: Dict[ObjectKeyType, model.HistoryDatasetCollectionAssociation]
+    hdcas_by_id: Dict[int, model.HistoryDatasetCollectionAssociation]
+    dces_by_key: Dict[ObjectKeyType, model.DatasetCollectionElement]
+    dces_by_id: Dict[int, model.DatasetCollectionElement]
+    lddas_by_key: Dict[ObjectKeyType, model.LibraryDatasetDatasetAssociation]
+    hda_copied_from_sinks: Dict[ObjectKeyType, ObjectKeyType]
+    hdca_copied_from_sinks: Dict[ObjectKeyType, ObjectKeyType]
+    jobs_by_key: Dict[ObjectKeyType, model.Job]
+    requires_hid: List[Union[model.HistoryDatasetAssociation, model.HistoryDatasetCollectionAssociation]]
 
     def __init__(self):
         self.libraries_by_key = {}
@@ -803,6 +810,28 @@ class ObjectImportTracker:
         self.jobs_by_key = {}
         self.requires_hid = []
 
+    def find_hda(self, input_key: ObjectKeyType) -> Optional[model.HistoryDatasetAssociation]:
+        hda = None
+        if input_key in self.hdas_by_key:
+            hda = self.hdas_by_key[input_key]
+        if input_key in self.hda_copied_from_sinks:
+            hda = self.hdas_by_key[self.hda_copied_from_sinks[input_key]]
+        return hda
+
+    def find_hdca(self, input_key: ObjectKeyType) -> Optional[model.HistoryDatasetCollectionAssociation]:
+        hdca = None
+        if input_key in self.hdcas_by_key:
+            hdca = self.hdcas_by_key[input_key]
+        if input_key in self.hdca_copied_from_sinks:
+            hdca = self.hdcas_by_key[self.hdca_copied_from_sinks[input_key]]
+        return hdca
+
+    def find_dce(self, input_key: ObjectKeyType) -> Optional[model.DatasetCollectionElement]:
+        dce = None
+        if input_key in self.dces_by_key:
+            dce = self.dces_by_key[input_key]
+        return dce
+
 
 def get_import_model_store_for_directory(archive_dir, **kwd):
     if not os.path.isdir(archive_dir):
@@ -815,6 +844,10 @@ def get_import_model_store_for_directory(archive_dir, **kwd):
 
 class BaseDirectoryImportModelStore(ModelImportStore):
 
+    @property
+    def file_source_root(self):
+        return self.archive_dir
+
     def defines_new_history(self):
         new_history_attributes = os.path.join(self.archive_dir, ATTRS_FILENAME_HISTORY)
         return os.path.exists(new_history_attributes)