Misc. edits/refactorings to session handling

lib/galaxy/app_unittest_utils/tools_support.py

@@ -133,12 +133,18 @@ def expunge_all(self):
     def query(self, clazz):
         return MockQuery(self.model_objects.get(clazz))
 
+    def get(self, clazz, id):
+        return self.query(clazz).get(id)
+
     def flush(self):
         self.flushed = True
 
     def add(self, object):
         self.created_objects.append(object)
 
+    def commit(self):
+        pass
+
 
 class MockQuery:
     def __init__(self, class_objects):

lib/galaxy/managers/history_contents.py

@@ -437,7 +437,7 @@ def _contained_id_map(self, id_list):
         query = (
             self._session()
             .query(component_class)
-            .filter(component_class.id.in_(id_list))
+            .filter(component_class.id.in_(id_list))  # type: ignore[attr-defined]
             .options(undefer(component_class._metadata))
             .options(joinedload(component_class.dataset).joinedload(model.Dataset.actions))
             .options(joinedload(component_class.tags))

lib/galaxy/model/__init__.py

@@ -4277,7 +4277,7 @@ def datatype_for_extension(extension, datatypes_registry=None) -> "Data":
     return ret
 
 
-class DatasetInstance(UsesCreateAndUpdateTime, _HasTable):
+class DatasetInstance(RepresentById, UsesCreateAndUpdateTime, _HasTable):
     """A base class for all 'dataset instances', HDAs, LDAs, etc"""
 
     states = Dataset.states

lib/galaxy/model/store/__init__.py

@@ -205,10 +205,10 @@ def commit(self) -> None:
     def flush(self) -> None:
         pass
 
-    def add(self, obj: Union[model.DatasetInstance, model.RepresentById]) -> None:
+    def add(self, obj: model.RepresentById) -> None:
         self.objects[obj.__class__][obj.id] = obj
 
-    def query(self, model_class: Type) -> Bunch:
+    def query(self, model_class: model.RepresentById) -> Bunch:
         def find(obj_id):
             return self.objects.get(model_class, {}).get(obj_id) or None
 
@@ -218,6 +218,9 @@ def filter_by(*args, **kwargs):
 
         return Bunch(find=find, get=find, filter_by=filter_by)
 
+    def get(self, model_class: model.RepresentById, primary_key: Any):  # patch for SQLAlchemy 2.0 compatibility
+        return self.query(model_class).get(primary_key)
+
 
 def replace_metadata_file(
     metadata: Dict[str, Any],
@@ -891,7 +894,7 @@ def materialize_elements(dc):
 
     def _attach_raw_id_if_editing(
         self,
-        obj: Union[model.DatasetInstance, model.RepresentById],
+        obj: model.RepresentById,
         attrs: Dict[str, Any],
     ) -> None:
         if self.sessionless and "id" in attrs and self.import_options.allow_edit:
@@ -1269,7 +1272,7 @@ def _import_implicit_collection_jobs(self, object_import_tracker: "ObjectImportT
 
             self._session_add(icj)
 
-    def _session_add(self, obj: Union[model.DatasetInstance, model.RepresentById]) -> None:
+    def _session_add(self, obj: model.RepresentById) -> None:
         self.sa_session.add(obj)
 
     def _flush(self) -> None:
@@ -1958,7 +1961,7 @@ def add(src, dest):
 
     def exported_key(
         self,
-        obj: Union[model.DatasetInstance, model.RepresentById],
+        obj: model.RepresentById,
     ) -> Union[str, int]:
         return self.serialization_options.get_identifier(self.security, obj)
 

lib/galaxy/tools/remote_tool_eval.py

@@ -60,7 +60,8 @@ def __init__(
         tool_data_table_manager: ToolDataTableManager,
         file_sources: ConfiguredFileSources,
     ):
-        self.model = Bunch(context=sa_session)
+        # For backward compatibility we need both context and session attributes that point to sa_session.
+        self.model = Bunch(context=sa_session, session=sa_session)
         self.config = tool_app_config
         self.datatypes_registry = datatypes_registry
         self.object_store = object_store

test/unit/app/jobs/test_job_wrapper.py

@@ -7,7 +7,10 @@
     Type,
 )
 
-from galaxy.app_unittest_utils.tools_support import UsesApp
+from galaxy.app_unittest_utils.tools_support import (
+    MockContext,
+    UsesApp,
+)
 from galaxy.jobs import (
     JobWrapper,
     TaskWrapper,
@@ -134,39 +137,6 @@ def url_to_destination(self):
         pass
 
 
-class MockContext:
-    def __init__(self, model_objects):
-        self.expunged_all = False
-        self.model_objects = model_objects
-        self.created_objects = []
-
-    def expunge_all(self):
-        self.expunged_all = True
-
-    def query(self, clazz):
-        return MockQuery(self.model_objects.get(clazz))
-
-    def flush(self):
-        pass
-
-    def commit(self):
-        pass
-
-    def add(self, object):
-        self.created_objects.append(object)
-
-
-class MockQuery:
-    def __init__(self, class_objects):
-        self.class_objects = class_objects
-
-    def filter_by(self, **kwds):
-        return Bunch(first=lambda: None)
-
-    def get(self, id):
-        return self.class_objects.get(id, None)
-
-
 class MockTool:
     def __init__(self, app):
         self.version_string_cmd = TEST_VERSION_COMMAND