Merge pull request #15094 from davelopez/fix_short_term_storage_api

Various fixes for Short Term Storage API

lib/galaxy/celery/_serialization.py

@@ -1,5 +1,6 @@
 import json
 from importlib import import_module
+from uuid import UUID
 
 from pydantic import BaseModel
 
@@ -23,6 +24,8 @@ class SchemaEncoder(json.JSONEncoder):
                 "__class__": fullname(obj),
                 "__object__": obj.dict(),
             }
+        if isinstance(obj, UUID):
+            return str(obj)
         else:
             return json.JSONEncoder.default(self, obj)
 

lib/galaxy/schema/schema.py

@@ -15,6 +15,7 @@ from typing import (
     Set,
     Union,
 )
+from uuid import UUID
 
 from pydantic import (
     AnyHttpUrl,
@@ -3191,7 +3192,7 @@ class AsyncTaskResultSummary(Model):
 
 
 class AsyncFile(Model):
-    storage_request_id: str
+    storage_request_id: UUID
     task: AsyncTaskResultSummary
 
 

lib/galaxy/schema/tasks.py

@@ -1,4 +1,5 @@
 from typing import Optional
+from uuid import UUID
 
 from pydantic import (
     BaseModel,
@@ -27,12 +28,12 @@ class SetupHistoryExportJob(BaseModel):
 
 
 class PrepareDatasetCollectionDownload(BaseModel):
-    short_term_storage_request_id: str
+    short_term_storage_request_id: UUID
     history_dataset_collection_association_id: int
 
 
 class GeneratePdfDownload(BaseModel):
-    short_term_storage_request_id: str
+    short_term_storage_request_id: UUID
     # basic markdown - Galaxy directives need to be processed before handing off to this task
     basic_markdown: str
     document_type: PdfDocumentType
@@ -47,14 +48,14 @@ class RequestUser(BaseModel):
 
 class GenerateHistoryDownload(StoreExportPayload):
     history_id: int
-    short_term_storage_request_id: str
+    short_term_storage_request_id: UUID
     user: RequestUser
 
 
 class GenerateHistoryContentDownload(StoreExportPayload):
     content_type: HistoryContentType
     content_id: int
-    short_term_storage_request_id: str
+    short_term_storage_request_id: UUID
     user: RequestUser
 
 
@@ -64,7 +65,7 @@ class BcoGenerationTaskParametersMixin(BcoGenerationParametersMixin):
 
 class GenerateInvocationDownload(StoreExportPayload, BcoGenerationTaskParametersMixin):
     invocation_id: int
-    short_term_storage_request_id: str
+    short_term_storage_request_id: UUID
     user: RequestUser
 
 

lib/galaxy/web/short_term_storage/__init__.py

@@ -12,12 +12,16 @@ from typing import (
     Optional,
     Union,
 )
-from uuid import uuid4
+from uuid import (
+    UUID,
+    uuid4,
+)
 
 from galaxy.exceptions import (
     InternalServerError,
     MessageException,
     NoContentException,
+    ObjectNotFound,
 )
 from galaxy.exceptions.error_codes import error_codes_by_int_code
 from galaxy.util import (
@@ -62,7 +66,7 @@ class ShortTermStorageTargetSecurity:
 
 @dataclass
 class ShortTermStorageTarget:
-    request_id: str  # uuid
+    request_id: UUID
     raw_path: str
 
     @property
@@ -144,7 +148,7 @@ class ShortTermStorageMonitor(metaclass=abc.ABCMeta):
         """Cleanup old requests."""
 
     @abc.abstractmethod
-    def recover_target(self, request_id: str) -> ShortTermStorageTarget:
+    def recover_target(self, request_id: UUID) -> ShortTermStorageTarget:
         """Return an existing ShortTermStorageTarget from a specified request_id."""
 
 
@@ -161,7 +165,7 @@ class ShortTermStorageManager(ShortTermStorageAllocator, ShortTermStorageMonitor
     ) -> ShortTermStorageTarget:
         if security is None:
             security = ShortTermStorageTargetSecurity()
-        request_id = str(uuid4())
+        request_id = uuid4()
         target_directory = self._directory(request_id)
         target = ShortTermStorageTarget(request_id=request_id, raw_path=str(target_directory / "target"))
         safe_makedirs(target_directory)
@@ -180,8 +184,7 @@ class ShortTermStorageManager(ShortTermStorageAllocator, ShortTermStorageMonitor
         self._store_metadata(target_directory, "request", request_info)
         return target
 
-    def recover_target(self, request_id: str) -> ShortTermStorageTarget:
-        assert is_uuid(request_id)  # secure the directory structure...
+    def recover_target(self, request_id: UUID) -> ShortTermStorageTarget:
         target_directory = self._directory(request_id)
         target = ShortTermStorageTarget(request_id=request_id, raw_path=str(target_directory / "target"))
         return target
@@ -244,18 +247,20 @@ class ShortTermStorageManager(ShortTermStorageAllocator, ShortTermStorageMonitor
 
     def _load_metadata(self, target_directory: Path, meta_name: str):
         meta_path = target_directory / f"{meta_name}.json"
+        if not meta_path.exists():
+            raise ObjectNotFound
         with open(meta_path) as f:
             return json.load(f)
 
-    def _directory(self, target: Union[str, ShortTermStorageTarget]) -> Path:
+    def _directory(self, target: Union[UUID, ShortTermStorageTarget]) -> Path:
         if isinstance(target, ShortTermStorageTarget):
             request_id = target.request_id
         else:
             request_id = target
-        relative_directory = directory_hash_id(request_id) + [request_id]
+        relative_directory = directory_hash_id(request_id) + [str(request_id)]
         return self._root.joinpath(*relative_directory)
 
-    def _cleanup_if_needed(self, request_id: str):
+    def _cleanup_if_needed(self, request_id: UUID):
         request_metadata = self._load_metadata(self._directory(request_id), "request")
         duration = request_metadata["duration"]
         creation_datetime_str = request_metadata["created"]
@@ -266,7 +271,7 @@ class ShortTermStorageManager(ShortTermStorageAllocator, ShortTermStorageMonitor
         if request_seconds > duration:
             self._delete(request_id)
 
-    def _delete(self, request_id: str):
+    def _delete(self, request_id: UUID):
         shutil.rmtree(self._directory(request_id))
 
     def cleanup(self):
@@ -274,7 +279,7 @@ class ShortTermStorageManager(ShortTermStorageAllocator, ShortTermStorageMonitor
             request_id = os.path.basename(directory)
             if not is_uuid(request_id):
                 continue
-            self._cleanup_if_needed(request_id)
+            self._cleanup_if_needed(UUID(request_id))
 
     @property
     def _root(self) -> Path:
@@ -282,7 +287,7 @@ class ShortTermStorageManager(ShortTermStorageAllocator, ShortTermStorageMonitor
 
 
 @contextlib.contextmanager
-def storage_context(short_term_storage_request_id: str, short_term_storage_monitor: ShortTermStorageMonitor):
+def storage_context(short_term_storage_request_id: UUID, short_term_storage_monitor: ShortTermStorageMonitor):
     target = short_term_storage_monitor.recover_target(short_term_storage_request_id)
     try:
         yield target

lib/galaxy/webapps/galaxy/api/short_term_storage.py

@@ -1,6 +1,8 @@
 """
 API operations around galaxy.web.short_term_storage infrastructure.
 """
+from uuid import UUID
+
 from galaxy.web.short_term_storage import (
     ShortTermStorageMonitor,
     ShortTermStorageServeCancelledInformation,
@@ -25,7 +27,7 @@ class FastAPIShortTermStorage:
         summary="Determine if specified storage request ID is ready for download.",
         response_description="Boolean indicating if the storage is ready.",
     )
-    def is_ready(self, storage_request_id: str) -> bool:
+    def is_ready(self, storage_request_id: UUID) -> bool:
         storage_target = self.short_term_storage_monitor.recover_target(storage_request_id)
         return self.short_term_storage_monitor.is_ready(storage_target)
 
@@ -43,7 +45,7 @@ class FastAPIShortTermStorage:
             },
         },
     )
-    def serve(self, storage_request_id: str):
+    def serve(self, storage_request_id: UUID):
         storage_target = self.short_term_storage_monitor.recover_target(storage_request_id)
         serve_info = self.short_term_storage_monitor.get_serve_info(storage_target)
         if isinstance(serve_info, ShortTermStorageServeCompletedInformation):

lib/galaxy_test/api/test_short_term_storage.py

@@ -0,0 +1,16 @@
+from uuid import uuid4
+
+from galaxy_test.base.api_asserts import assert_status_code_is
+from ._framework import ApiTestCase
+
+
+class TestShortTermStorageApi(ApiTestCase):
+    def test_invalid_request_id_returns_400(self):
+        invalid_uuid = "invalid_uuid"
+        response = self._get(f"short_term_storage/{invalid_uuid}")
+        assert_status_code_is(response, 400)
+
+    def test_non_existent_request_id_returns_404(self):
+        valid_uuid = uuid4()
+        response = self._get(f"short_term_storage/{valid_uuid}")
+        assert_status_code_is(response, 404)

lib/galaxy_test/base/populators.py

@@ -64,6 +64,7 @@ from typing import (
     Tuple,
     Union,
 )
+from uuid import UUID
 
 import cwltest.utils
 import requests
@@ -1209,7 +1210,7 @@ class BaseDatasetPopulator(BasePopulator):
         storage_request_id = self.assert_download_request_ok(download_request_response)
         return self.wait_on_download_request(storage_request_id)
 
-    def assert_download_request_ok(self, download_request_response: Response) -> str:
+    def assert_download_request_ok(self, download_request_response: Response) -> UUID:
         """Assert response is valid and okay and extract storage request ID."""
         api_asserts.assert_status_code_is(download_request_response, 200)
         download_async = download_request_response.json()
@@ -1217,7 +1218,7 @@ class BaseDatasetPopulator(BasePopulator):
         storage_request_id = download_async["storage_request_id"]
         return storage_request_id
 
-    def wait_for_download_ready(self, storage_request_id: str):
+    def wait_for_download_ready(self, storage_request_id: UUID):
         def is_ready():
             is_ready_response = self._get(f"short_term_storage/{storage_request_id}/ready")
             is_ready_response.raise_for_status()
@@ -1244,7 +1245,7 @@ class BaseDatasetPopulator(BasePopulator):
         wait_on(is_ready, "waiting for task to complete")
         return state() == "SUCCESS"
 
-    def wait_on_download_request(self, storage_request_id: str) -> Response:
+    def wait_on_download_request(self, storage_request_id: UUID) -> Response:
         self.wait_for_download_ready(storage_request_id)
         download_contents_response = self._get(f"short_term_storage/{storage_request_id}")
         download_contents_response.raise_for_status()

test/unit/web_framework/test_short_term_storage.py

@@ -1,6 +1,11 @@
 import time
 
-from galaxy.exceptions import MessageException
+import pytest
+
+from galaxy.exceptions import (
+    MessageException,
+    ObjectNotFound,
+)
 from galaxy.web.short_term_storage import (
     ShortTermStorageConfiguration,
     ShortTermStorageManager,
@@ -121,3 +126,24 @@ def test_cleanup_no_op_if_duration_not_reached(tmpdir):
     time.sleep(TEST_SLEEP_DURATION)
     manager.cleanup()
     assert short_term_storage_target.path.exists()
+
+
+def test_serve_non_existent_raises_object_not_found(tmpdir):
+    config = ShortTermStorageConfiguration(
+        short_term_storage_directory=tmpdir,
+        maximum_storage_duration=TEST_SLEEP_DURATION / 10,
+    )
+    manager = ShortTermStorageManager(config=config)
+    short_term_storage_target = manager.new_target(
+        TEST_FILENAME,
+        TEST_MIME_TYPE,
+    )
+    short_term_storage_target.path.touch()
+    assert short_term_storage_target.path.exists()
+    serve_info = manager.get_serve_info(short_term_storage_target)
+    assert serve_info
+    time.sleep(TEST_SLEEP_DURATION)
+    manager.cleanup()
+    assert not short_term_storage_target.path.exists()
+    with pytest.raises(ObjectNotFound):
+        manager.get_serve_info(short_term_storage_target)