Track direct downloads to reuse the STS export

client/src/components/Common/ExportRecordDetails.vue

@@ -18,7 +18,7 @@ const props = defineProps({
     },
 });
 
-const emit = defineEmits(["onReimport"]);
+const emit = defineEmits(["onReimport", "onDownload"]);
 
 const title = computed(() => (props.record.isReady ? `Exported` : `Export started`));
 const elapsedTime = computed(() => formatDistanceToNow(parseISO(`${props.record.date}Z`), { addSuffix: true }));
@@ -39,6 +39,10 @@ const preparingMessage = computed(
 function reimportObject() {
     emit("onReimport", props.record);
 }
+
+function downloadObject() {
+    emit("onDownload", props.record);
+}
 </script>
 
 <template>
@@ -60,6 +64,9 @@ function reimportObject() {
                 <p class="mt-3">
                     {{ readyMessage }}
                 </p>
+                <b-button v-if="props.record.canDownload" variant="primary" @click="downloadObject">
+                    Download
+                </b-button>
                 <b-button v-if="props.record.canReimport" variant="primary" @click="reimportObject">
                     Reimport
                 </b-button>

client/src/components/Common/models/exportRecordModel.js

@@ -27,12 +27,23 @@ export class ExportRecordModel {
         return this._data.task_uuid;
     }
 
-    // import_uri doesn't work for downloads
     get canReimport() {
-        return this._data?.export_metadata?.result_data?.import_uri !== undefined;
+        return this.isReady && !!this._data?.export_metadata?.result_data?.target_uri;
     }
 
     get importLink() {
-        return this._data?.export_metadata?.result_data?.import_uri;
+        return this._data?.export_metadata?.result_data?.target_uri;
+    }
+
+    get isStsDownload() {
+        return !!this._data?.export_metadata?.result_data?.short_term_storage_request_id;
+    }
+
+    get stsDownloadId() {
+        return this._data?.export_metadata?.result_data?.short_term_storage_request_id;
+    }
+
+    get canDownload() {
+        return this.isReady && this.isStsDownload;
     }
 }

client/src/components/History/Export/HistoryExport.vue

@@ -13,7 +13,7 @@ const service = new HistoryExportService();
 
 const { isRunning: isExportTaskRunning, waitForTask } = useTaskMonitor();
 const { hasWritable: hasWritableFileSources } = useFileSources();
-const { isPreparing: isPreparingDownload, downloadHistory } = useShortTermStorage();
+const { isPreparing: isPreparingDownload, downloadHistory, downloadObjectByRequestId } = useShortTermStorage();
 
 const props = defineProps({
     historyId: {
@@ -70,11 +70,22 @@ async function exportToFileSource(exportDirectory, fileName) {
 }
 
 async function prepareDownload() {
-    downloadHistory(props.historyId, { pollDelayInMs: 3000, exportParams: EXPORT_PARAMS });
-    // updateExports();
+    if (latestExportRecord.value?.isStsDownload && latestExportRecord.value.isUpToDate) {
+        console.debug("Existing STS download found");
+        downloadObjectByRequestId(latestExportRecord.value.stsDownloadId);
+        return;
+    }
+    await downloadHistory(props.historyId, { pollDelayInMs: 3000, exportParams: EXPORT_PARAMS });
+    updateExports();
+}
+
+function downloadFromRecord(record) {
+    if (record.isStsDownload) {
+        downloadObjectByRequestId(record.stsDownloadId);
+    }
 }
 
-function reimportHistoryFromRecord(record) {
+function reimportFromRecord(record) {
     return service.reimportHistoryFromRecord(record);
 }
 </script>
@@ -111,7 +122,8 @@ function reimportHistoryFromRecord(record) {
             :record="latestExportRecord"
             object-type="history"
             class="mt-3"
-            @onReimport="reimportHistoryFromRecord" />
+            @onDownload="downloadFromRecord"
+            @onReimport="reimportFromRecord" />
         <b-alert v-else-if="errorMessage" variant="danger" class="mt-3" show>
             {{ errorMessage }}
         </b-alert>

client/src/composables/shortTermStorage.js

@@ -21,11 +21,11 @@ export function useShortTermStorage() {
 
     const isPreparing = ref(false);
 
-    function downloadHistory(historyId, options = DEFAULT_OPTIONS) {
+    async function downloadHistory(historyId, options = DEFAULT_OPTIONS) {
         return prepareObjectDownload(historyId, "histories", options);
     }
 
-    function downloadWorkflowInvocation(invocationId, options = DEFAULT_OPTIONS) {
+    async function downloadWorkflowInvocation(invocationId, options = DEFAULT_OPTIONS) {
         return prepareObjectDownload(invocationId, "invocations", options);
     }
 
@@ -34,7 +34,7 @@ export function useShortTermStorage() {
         window.location.assign(url);
     }
 
-    function prepareObjectDownload(object_id, object_api, options = DEFAULT_OPTIONS) {
+    async function prepareObjectDownload(object_id, object_api, options = DEFAULT_OPTIONS) {
         const finalOptions = Object.assign(DEFAULT_OPTIONS, options);
         resetTimeout();
         isPreparing.value = true;
@@ -46,7 +46,9 @@ export function useShortTermStorage() {
             include_deleted: finalOptions.exportParams.includeDeleted,
             include_hidden: finalOptions.exportParams.includeHidden,
         };
-        axios.post(url, exportParams).then(handleInitialize).catch(handleError);
+
+        const response = await axios.post(url, exportParams).catch(handleError);
+        handleInitialize(response);
     }
 
     function handleInitialize(response) {

lib/galaxy/managers/model_stores.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union
 
 from galaxy import model
 from galaxy.exceptions import RequestParameterInvalidException
@@ -20,6 +20,7 @@
     ExportObjectType,
     HistoryContentType,
     WriteStoreToPayload,
+    ShortTermStoreExportPayload,
 )
 from galaxy.schema.tasks import (
     BcoGenerationTaskParametersMixin,
@@ -86,6 +87,7 @@ def prepare_history_download(self, request: GenerateHistoryDownload):
                 short_term_storage_target.path
             ) as export_store:
                 export_store.export_history(history, include_hidden=include_hidden, include_deleted=include_deleted)
+            self.set_history_export_metadata(request)
 
     def prepare_history_content_download(self, request: GenerateHistoryContentDownload):
         model_store_format = request.model_store_format
@@ -177,16 +179,22 @@ def write_history_to(self, request: WriteHistoryTo):
             )
             self.set_history_export_metadata(request)
 
-    def set_history_export_metadata(self, request: WriteHistoryTo):
+    def set_history_export_metadata(self, request: Union[WriteHistoryTo, GenerateHistoryDownload]):
         if request.export_association_id:
+            request_dict = request.dict()
+            request_payload = (
+                WriteStoreToPayload(**request_dict)
+                if isinstance(request, WriteHistoryTo)
+                else ShortTermStoreExportPayload(**request_dict)
+            )
             export_metadata = ExportObjectMetadata(
                 request_data=ExportObjectRequestMetadata(
                     object_id=request.history_id,
                     object_type=ExportObjectType.HISTORY,
                     user_id=request.user.user_id,
-                    payload=WriteStoreToPayload(**request.dict()),
+                    payload=request_payload,
                 ),
-                result_data=ExportObjectResultMetadata(import_uri=request.target_uri),
+                result_data=ExportObjectResultMetadata(**request_dict),
             )
             self._export_tracker.set_export_association_metadata(request.export_association_id, export_metadata)
 

lib/galaxy/schema/schema.py

@@ -1428,11 +1428,12 @@ class ExportObjectRequestMetadata(Model):
     object_id: EncodedDatabaseIdField
     object_type: ExportObjectType
     user_id: Optional[EncodedDatabaseIdField]
-    payload: WriteStoreToPayload
+    payload: Union[WriteStoreToPayload, ShortTermStoreExportPayload]
 
 
 class ExportObjectResultMetadata(Model):
-    import_uri: str
+    target_uri: Optional[str]
+    short_term_storage_request_id: Optional[str]
 
 
 class ExportObjectMetadata(Model):

lib/galaxy/webapps/galaxy/services/histories.py

@@ -337,14 +337,19 @@ def prepare_download(
             history.name,
             payload.model_store_format,
         )
+        export_association = self.history_export_manager.create_export_association(history.id)
         request = GenerateHistoryDownload(
             history_id=history.id,
             short_term_storage_request_id=short_term_storage_target.request_id,
             user=trans.async_request_user,
+            export_association_id=export_association.id,
             **payload.dict(),
         )
         result = prepare_history_download.delay(request=request)
-        return AsyncFile(storage_request_id=short_term_storage_target.request_id, task=async_task_summary(result))
+        task_summary = async_task_summary(result)
+        export_association.task_uuid = task_summary.id
+        trans.sa_session.flush()
+        return AsyncFile(storage_request_id=short_term_storage_target.request_id, task=task_summary)
 
     def write_store(
         self, trans: ProvidesHistoryContext, history_id: DecodedDatabaseIdField, payload: WriteStoreToPayload
@@ -358,11 +363,10 @@ def write_store(
             **payload.dict(),
         )
         result = write_history_to.delay(request=request)
-        summary = async_task_summary(result)
-        # TODO: will this race?
-        export_association.task_uuid = summary.id
+        task_summary = async_task_summary(result)
+        export_association.task_uuid = task_summary.id
         trans.sa_session.flush()
-        return summary
+        return task_summary
 
     def update(
         self,