Merge pull request #109 from i-VRESSE/archive-filename

Add filename query search param to archive route

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bartender"
-version = "0.5.2"
+version = "0.5.3"
 description = "Job middleware for i-VRESSE"
 authors = [
 

src/bartender/web/api/job/views.py

@@ -370,12 +370,13 @@ def _remove_archive(filename: str) -> None:
     },
     response_class=FileResponse,
 )
-async def retrieve_job_directory_as_archive(
+async def retrieve_job_directory_as_archive(  # noqa: WPS211
     job_dir: CurrentCompletedJobDir,
     background_tasks: BackgroundTasks,
     archive_format: ArchiveFormat = ".zip",
     exclude: Optional[list[str]] = Query(default=None),
     exclude_dirs: Optional[list[str]] = Query(default=None),
+    filename: Optional[str] = Query(default=None),
     # Note: also tried to with include (filter & filter_dirs) but that can be
     # unintuitive. e.g. include_dirs=['output'] doesn't return subdirs of
     # /output that are not also called output. Might improve when globs will be
@@ -391,6 +392,8 @@ async def retrieve_job_directory_as_archive(
             '.tar.xz', '.tar.gz', '.tar.bz2'
         exclude: list of filename patterns that should be excluded from archive.
         exclude_dirs: list of directory patterns that should be excluded from archive.
+        filename: Name of the archive file to be returned.
+            If not provided, uses id of the job.
 
     Returns:
         FileResponse: Archive containing the content of job_dir
@@ -402,6 +405,8 @@ async def retrieve_job_directory_as_archive(
     background_tasks.add_task(_remove_archive, archive_fn)
 
     return_fn = Path(archive_fn).name
+    if filename:
+        return_fn = filename
     return FileResponse(archive_fn, filename=return_fn)
 
 
@@ -413,6 +418,7 @@ async def retrieve_job_subdirectory_as_archive(  # noqa: WPS211
     archive_format: ArchiveFormat = ".zip",
     exclude: Optional[list[str]] = Query(default=None),
     exclude_dirs: Optional[list[str]] = Query(default=None),
+    filename: Optional[str] = Query(default=None),
 ) -> FileResponse:
     """Download job output as archive.
 
@@ -424,6 +430,8 @@ async def retrieve_job_subdirectory_as_archive(  # noqa: WPS211
             '.tar', '.tar.xz', '.tar.gz', '.tar.bz2'
         exclude: list of filename patterns that should be excluded from archive.
         exclude_dirs: list of directory patterns that should be excluded from archive.
+        filename: Name of the archive file to be returned.
+            If not provided, uses id of the job.
 
     Returns:
         FileResponse: Archive containing the output of job_dir
@@ -436,6 +444,7 @@ async def retrieve_job_subdirectory_as_archive(  # noqa: WPS211
         archive_format=archive_format,
         exclude=exclude,
         exclude_dirs=exclude_dirs,
+        filename=filename,
     )
 
 

tests/web/test_job.py

@@ -758,6 +758,45 @@ async def test_job_directory_as_archive(
     assert stdout == "this is stdout"
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize(
+    "archive_format",
+    [".zip"],
+)
+async def test_job_directory_as_named_archive(
+    fastapi_app: FastAPI,
+    client: AsyncClient,
+    auth_headers: Dict[str, str],
+    mock_ok_job: int,
+    archive_format: str,
+) -> None:
+    url = (
+        fastapi_app.url_path_for(
+            "retrieve_job_directory_as_archive",
+            jobid=mock_ok_job,
+        )
+        + f"?archive_format={archive_format}&filename=foo.zip"
+    )
+    response = await client.get(url, headers=auth_headers)
+
+    expected_content_type = (
+        "application/zip" if archive_format == ".zip" else "application/x-tar"
+    )
+    expected_content_disposition = f'attachment; filename="foo{archive_format}"'
+
+    assert response.status_code == status.HTTP_200_OK
+    assert response.headers["content-type"] == expected_content_type
+    assert response.headers["content-disposition"] == expected_content_disposition
+
+    fs = ZipFS if archive_format == ".zip" else TarFS
+
+    with io.BytesIO(response.content) as responsefile:
+        with fs(responsefile) as archive:
+            stdout = archive.readtext("stdout.txt")
+
+    assert stdout == "this is stdout"
+
+
 @pytest.mark.anyio
 async def test_job_subdirectory_as_archive(
     fastapi_app: FastAPI,
@@ -785,6 +824,34 @@ async def test_job_subdirectory_as_archive(
     assert stdout == "hi from output dir"
 
 
+@pytest.mark.anyio
+async def test_job_subdirectory_as_named_archive(
+    fastapi_app: FastAPI,
+    client: AsyncClient,
+    auth_headers: Dict[str, str],
+    mock_ok_job: int,
+) -> None:
+    url = (
+        fastapi_app.url_path_for(
+            "retrieve_job_subdirectory_as_archive",
+            jobid=mock_ok_job,
+            path="output",
+        )
+        + "?filename=bar.zip"
+    )
+    response = await client.get(url, headers=auth_headers)
+
+    assert response.status_code == status.HTTP_200_OK
+    assert response.headers["content-type"] == "application/zip"
+    assert response.headers["content-disposition"] == 'attachment; filename="bar.zip"'
+
+    with io.BytesIO(response.content) as responsefile:
+        with ZipFS(responsefile) as archive:
+            stdout = archive.readtext("readme.txt")
+
+    assert stdout == "hi from output dir"
+
+
 @pytest.fixture
 def demo_interactive_application(
     demo_config: Config,