-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Drop jobdir/meta file + refactor bartender.filesystem modules to stag…
…ing and walk_dir module
- Loading branch information
1 parent
83f2384
commit 129b068
Showing
10 changed files
with
200 additions
and
226 deletions.
There are no files selected for viewing
This file was deleted.
Oops, something went wrong.
This file was deleted.
Oops, something went wrong.
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,115 @@ | ||
from pathlib import Path | ||
from shutil import unpack_archive | ||
from typing import Literal, Optional | ||
|
||
from aiofiles import open | ||
from aiofiles.os import mkdir, remove | ||
from fastapi import HTTPException, UploadFile | ||
from starlette import status | ||
|
||
from bartender.async_utils import async_wrap | ||
from bartender.config import ApplicatonConfiguration | ||
|
||
CHUNK_SIZE = 1024 * 1024 # 1Mb | ||
|
||
|
||
class UnsupportedContentTypeError(Exception): | ||
"""When content type is unsupported.""" | ||
|
||
|
||
def has_needed_files( | ||
application: ApplicatonConfiguration, | ||
job_dir: Path, | ||
) -> Literal[True]: | ||
"""Check if files required by application are present in job directory. | ||
Args: | ||
application: Name of application to check config file for. | ||
job_dir: In which directory to look. | ||
Raises: | ||
IndexError: When one or more needed files can not be found | ||
Returns: | ||
True when found or no files where needed. | ||
""" | ||
missing_files = [] | ||
for needed_file in application.upload_needs: | ||
file = job_dir / needed_file | ||
file_exists = file.exists() and file.is_file() | ||
if not file_exists: | ||
missing_files.append(needed_file) | ||
if missing_files: | ||
raise IndexError( | ||
f"Application requires files {missing_files}, " | ||
"but where not found in uploaded zip archive", | ||
) | ||
return True | ||
|
||
|
||
async def create_job_dir(job_id: int, job_root_dir: Path) -> Path: | ||
"""Create job directory. | ||
Args: | ||
job_id: id of the job. | ||
job_root_dir: Root directory for all jobs. | ||
Raises: | ||
HTTPException: When job directory could not be made. | ||
Returns: | ||
Directory of job. | ||
""" | ||
job_dir: Path = job_root_dir / str(job_id) | ||
|
||
try: | ||
await mkdir(job_dir) | ||
except FileExistsError as exc: | ||
raise HTTPException( | ||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, | ||
detail="Failed to create directory for job.", | ||
) from exc | ||
return job_dir | ||
|
||
|
||
def _is_valid_content_type(content_type: Optional[str]) -> bool: | ||
supported_upload_content_types = { | ||
"application/zip", | ||
"application/x-zip-compressed", | ||
} # TODO add support for other formats like tar.gz, tar.bz2, .7z? | ||
if content_type not in supported_upload_content_types: | ||
raise UnsupportedContentTypeError( | ||
f"Unable to stage job input wrong mime type {content_type}, " | ||
+ f"supported are {supported_upload_content_types}", | ||
) | ||
return True | ||
|
||
|
||
async def unpack_upload( | ||
job_dir: Path, | ||
archive: UploadFile, | ||
dest_fn: str = "archive.zip", | ||
) -> None: | ||
"""Unpack archive file to job id directory. | ||
Args: | ||
job_dir: Where to put archive file. | ||
archive: The archive file with async read method. | ||
dest_fn: Temporary archive filename. | ||
""" | ||
_is_valid_content_type(archive.content_type) | ||
|
||
# Copy archive to disk | ||
job_archive = job_dir / dest_fn | ||
# If archive contains | ||
async with open(job_archive, "wb") as out_file: | ||
while content := await archive.read(CHUNK_SIZE): | ||
if isinstance(content, str): | ||
break # type narrowing for mypy, content is always bytes | ||
await out_file.write(content) | ||
|
||
if archive.content_type in {"application/zip", "application/x-zip-compressed"}: | ||
await async_wrap(unpack_archive)(job_archive, extract_dir=job_dir, format="zip") | ||
# TODO what happens when archive contains archive.zip, will it overwrite itself? | ||
|
||
await remove(job_archive) # no longer needed? |
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,78 @@ | ||
from pathlib import Path | ||
|
||
import pytest | ||
|
||
from bartender.config import ApplicatonConfiguration | ||
from bartender.staging import create_job_dir, has_needed_files | ||
|
||
|
||
@pytest.mark.anyio | ||
async def test_create_job_dir(job_root_dir: Path) -> None: | ||
"""Test the assembly the job.""" | ||
job_id = 1 | ||
|
||
await create_job_dir(job_id, job_root_dir) | ||
|
||
job_dir = job_root_dir / str(job_id) | ||
assert job_dir.exists() | ||
|
||
|
||
class TestHasNeededFiles: | ||
@pytest.fixture | ||
def job_dir(self, tmp_path: Path) -> Path: | ||
return tmp_path | ||
|
||
@pytest.fixture | ||
def application(self) -> ApplicatonConfiguration: | ||
return ApplicatonConfiguration( | ||
command_template="wc config.ini data.csv", | ||
upload_needs=["config.ini", "data.csv"], | ||
) | ||
|
||
def test_has_needed_files_with_existing_files( | ||
self, | ||
job_dir: Path, | ||
application: ApplicatonConfiguration, | ||
) -> None: | ||
# Create the needed files | ||
(job_dir / "config.ini").touch() | ||
(job_dir / "data.csv").touch() | ||
|
||
# Check if the files exist | ||
result = has_needed_files(application, job_dir) | ||
|
||
# Assert that the function returns True | ||
assert result is True | ||
|
||
def test_has_needed_files_with_missing_files( | ||
self, | ||
job_dir: Path, | ||
application: ApplicatonConfiguration, | ||
) -> None: | ||
# Check if the files are missing | ||
with pytest.raises(IndexError): | ||
has_needed_files(application, job_dir) | ||
|
||
def test_has_needed_files_with_partial_missing_files( | ||
self, | ||
job_dir: Path, | ||
application: ApplicatonConfiguration, | ||
) -> None: | ||
# Create one of the needed files | ||
(job_dir / "config.ini").touch() | ||
|
||
# Check if the files are missing | ||
with pytest.raises(IndexError): | ||
has_needed_files(application, job_dir) | ||
|
||
def test_has_needed_files_with_no_files_needed( | ||
self, | ||
job_dir: Path, | ||
application: ApplicatonConfiguration, | ||
) -> None: | ||
# Remove the upload_needs list | ||
application.upload_needs = [] | ||
|
||
# Check if the function returns True | ||
result = has_needed_files(application, job_dir) | ||
assert result is True |
Oops, something went wrong.