From a8191d4b8a4bcf1b7d0cd303a42838b608026b2c Mon Sep 17 00:00:00 2001
From: mvdbeek <m.vandenbeek@gmail.com>
Date: Sun, 25 Oct 2020 16:48:59 +0100
Subject: [PATCH] Move one more flush out of dicovering outputs

---
 lib/galaxy/model/__init__.py       | 5 +++++
 lib/galaxy/model/store/discover.py | 1 -
 lib/galaxy/objectstore/__init__.py | 6 +++++-
 test/unit/jobs/test_job_context.py | 3 ++-
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/lib/galaxy/model/__init__.py b/lib/galaxy/model/__init__.py
index db92f3f64e8..bac0d7dd509 100644
--- a/lib/galaxy/model/__init__.py
+++ b/lib/galaxy/model/__init__.py
@@ -2307,6 +2307,11 @@ class StorableObject:
         else:
             self.uuid = UUID(str(uuid))
 
+    def flush(self):
+        sa_session = object_session(self)
+        if sa_session:
+            sa_session.flush()
+
 
 class Dataset(StorableObject, RepresentById):
     states = Bunch(NEW='new',
diff --git a/lib/galaxy/model/store/discover.py b/lib/galaxy/model/store/discover.py
index 263aa8d794a..bfee249db29 100644
--- a/lib/galaxy/model/store/discover.py
+++ b/lib/galaxy/model/store/discover.py
@@ -283,7 +283,6 @@ class ModelPersistenceContext(metaclass=abc.ABCMeta):
             association_name = f'__new_primary_file_{name}|{element_identifier_str}__'
             self.add_output_dataset_association(association_name, dataset)
 
-        self.flush()
         self.update_object_store_with_datasets(datasets=element_datasets['datasets'], paths=element_datasets['paths'], extra_files=element_datasets['extra_files'])
         add_datasets_timer = ExecutionTimer()
         self.add_datasets_to_history(element_datasets['datasets'])
diff --git a/lib/galaxy/objectstore/__init__.py b/lib/galaxy/objectstore/__init__.py
index 2166b2a8120..c9cfaecd6fc 100644
--- a/lib/galaxy/objectstore/__init__.py
+++ b/lib/galaxy/objectstore/__init__.py
@@ -256,7 +256,11 @@ class BaseObjectStore(ObjectStore):
 
     def _get_object_id(self, obj):
         if hasattr(obj, self.store_by):
-            return getattr(obj, self.store_by)
+            obj_id = getattr(obj, self.store_by)
+            if obj_id is None:
+                obj.flush()
+                return obj.id
+            return obj_id
         else:
             # job's don't have uuids, so always use ID in this case when creating
             # job working directories.
diff --git a/test/unit/jobs/test_job_context.py b/test/unit/jobs/test_job_context.py
index faf0ed992f8..6bbacc62331 100644
--- a/test/unit/jobs/test_job_context.py
+++ b/test/unit/jobs/test_job_context.py
@@ -81,6 +81,7 @@ def test_job_context_discover_outputs_flushes_once(mocker):
         final_job_state=job_context.final_job_state,
     )
     collection_builder.populate()
-    assert spy.call_count == 1
+    assert spy.call_count == 0
+    sa_session.flush()
     assert len(collection.dataset_instances) == 10
     assert collection.dataset_instances[0].dataset.file_size == 1