Skip to content

Commit

Permalink
Add some test coverage and clean up code
Browse files Browse the repository at this point in the history
I removed unused fixtures (particularly supporting the obsolete "controller"
header).

Also slightly rearranged some code in upload: I removed a new exception handler
that's not really necessary and realized that some post-upload checks were
unnecessarily inside the `with` block.
  • Loading branch information
dbutenhof committed Mar 21, 2023
1 parent c3f94a1 commit 333e6c5
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 105 deletions.
88 changes: 42 additions & 46 deletions lib/pbench/server/api/resources/upload_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,10 +248,10 @@ def _put(self, args: ApiParams, request: Request, context: ApiContext) -> Respon
try:
tmp_dir = self.temporary / md5sum
tmp_dir.mkdir()
except FileExistsError as e:
except FileExistsError:
raise CleanupTime(
HTTPStatus.INTERNAL_SERVER_ERROR,
f"Temporary upload directory {tmp_dir} already exists: {e}",
HTTPStatus.CONFLICT,
"Temporary upload directory already exists",
)
upload_dir = tmp_dir
tar_full_path = upload_dir / filename
Expand Down Expand Up @@ -326,65 +326,61 @@ def _put(self, args: ApiParams, request: Request, context: ApiContext) -> Respon
# error recovery.
recovery.add(tar_full_path.unlink)

# Open for exclusive WRITE in BINARY mode
with tar_full_path.open(mode="xb") as ofp:
hash_md5 = hashlib.md5()
# NOTE: We know that the MD5 is unique at this point; so even if
# two tarballs with the same name are uploaded concurrently, by
# writing into a temporary directory named for the MD5 we're
# assured that they can't conflict.
try:
with tar_full_path.open(mode="wb") as ofp:
hash_md5 = hashlib.md5()

try:
while True:
chunk = request.stream.read(self.CHUNK_SIZE)
bytes_received += len(chunk)
if len(chunk) == 0 or bytes_received > content_length:
break
ofp.write(chunk)
hash_md5.update(chunk)
except FileExistsError:
except OSError as exc:
if exc.errno == errno.ENOSPC:
raise CleanupTime(
HTTPStatus.CONFLICT,
f"A tarball named {dataset_name!r} is already being uploaded",
HTTPStatus.INSUFFICIENT_STORAGE,
f"Out of space on {tar_full_path.root}",
)
except OSError as exc:
if exc.errno == errno.ENOSPC:
raise CleanupTime(
HTTPStatus.INSUFFICIENT_STORAGE,
f"Out of space on {tar_full_path.root}",
)
else:
raise CleanupTime(
HTTPStatus.INTERNAL_SERVER_ERROR,
f"Unexpected error {exc.errno} encountered during file upload",
)
except Exception:
else:
raise CleanupTime(
HTTPStatus.INTERNAL_SERVER_ERROR,
"Unexpected error encountered during file upload",
)

if bytes_received != content_length:
raise CleanupTime(
HTTPStatus.BAD_REQUEST,
f"Expected {content_length} bytes but received {bytes_received} bytes",
)
elif hash_md5.hexdigest() != md5sum:
raise CleanupTime(
HTTPStatus.BAD_REQUEST,
f"MD5 checksum {hash_md5.hexdigest()} does not match expected {md5sum}",
f"Unexpected error {exc.errno} encountered during file upload",
)
except Exception:
raise CleanupTime(
HTTPStatus.INTERNAL_SERVER_ERROR,
"Unexpected error encountered during file upload",
)

# First write the .md5
current_app.logger.info(
"Creating MD5 file {}: {}", md5_full_path, md5sum
if bytes_received != content_length:
raise CleanupTime(
HTTPStatus.BAD_REQUEST,
f"Expected {content_length} bytes but received {bytes_received} bytes",
)
elif hash_md5.hexdigest() != md5sum:
raise CleanupTime(
HTTPStatus.BAD_REQUEST,
f"MD5 checksum {hash_md5.hexdigest()} does not match expected {md5sum}",
)

# From this point attempt to remove the MD5 file on error exit
recovery.add(md5_full_path.unlink)
try:
md5_full_path.write_text(f"{md5sum} {filename}\n")
except Exception:
raise CleanupTime(
HTTPStatus.INTERNAL_SERVER_ERROR,
f"Failed to write .md5 file '{md5_full_path}'",
)
# First write the .md5
current_app.logger.info("Creating MD5 file {}: {}", md5_full_path, md5sum)

# From this point attempt to remove the MD5 file on error exit
recovery.add(md5_full_path.unlink)
try:
md5_full_path.write_text(f"{md5sum} {filename}\n")
except Exception:
raise CleanupTime(
HTTPStatus.INTERNAL_SERVER_ERROR,
f"Failed to write .md5 file '{md5_full_path}'",
)

# Create a cache manager object
try:
Expand Down
Loading

0 comments on commit 333e6c5

Please sign in to comment.