From 9646f6367ece7a99833469d26a5aa44de9e5aa9b Mon Sep 17 00:00:00 2001 From: c-bata Date: Fri, 2 Jun 2023 18:02:55 +0900 Subject: [PATCH] Add ArtifactNotFound exception --- docs/api.rst | 2 + optuna_dashboard/artifact/boto3.py | 22 ++++++++++- optuna_dashboard/artifact/exceptions.py | 9 +++++ optuna_dashboard/artifact/file_system.py | 16 ++++++-- optuna_dashboard/artifact/protocol.py | 45 ++++++++++++++++++++++- python_tests/test_boto3_artifact.py | 6 +++ python_tests/test_file_system_artifact.py | 8 ++++ 7 files changed, 101 insertions(+), 7 deletions(-) create mode 100644 optuna_dashboard/artifact/exceptions.py diff --git a/docs/api.rst b/docs/api.rst index 061e93fc..a286b2c7 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -40,3 +40,5 @@ Artifact optuna_dashboard.artifact.upload_artifact optuna_dashboard.artifact.file_system.FileSystemBackend optuna_dashboard.artifact.boto3.Boto3Backend + optuna_dashboard.artifact.protocol.ArtifactBackend + optuna_dashboard.artifact.exceptions.ArtifactNotFound diff --git a/optuna_dashboard/artifact/boto3.py b/optuna_dashboard/artifact/boto3.py index b248d5dc..e3da2ef0 100644 --- a/optuna_dashboard/artifact/boto3.py +++ b/optuna_dashboard/artifact/boto3.py @@ -5,6 +5,8 @@ import shutil from typing import TYPE_CHECKING import boto3 +from botocore.exceptions import ClientError +from optuna_dashboard.artifact.exceptions import ArtifactNotFound if TYPE_CHECKING: @@ -41,7 +43,12 @@ class Boto3Backend: self.client = client or boto3.client("s3") def open(self, artifact_id: str) -> BinaryIO: - obj = self.client.get_object(Bucket=self.bucket, Key=artifact_id) + try: + obj = self.client.get_object(Bucket=self.bucket, Key=artifact_id) + except ClientError as e: + if _is_not_found_error(e): + raise ArtifactNotFound("not found") from e + raise body = obj.get("Body") assert body is not None return body # type: ignore @@ -58,7 +65,18 @@ class Boto3Backend: self.client.upload_fileobj(buf, self.bucket, artifact_id) def remove(self, artifact_id: str) -> None: - self.client.delete_object(Bucket=self.bucket, Key=artifact_id) + try: + self.client.delete_object(Bucket=self.bucket, Key=artifact_id) + except ClientError as e: + if _is_not_found_error(e): + raise ArtifactNotFound("not found") from e + raise + + +def _is_not_found_error(e: ClientError) -> bool: + error_code = e.response.get("Error", {}).get("Code") + http_status_code = e.response.get("ResponseMetadata", {}).get("HTTPStatusCode") + return error_code == "NoSuchKey" or http_status_code == 404 def _is_file_like_obj(obj: SupportsRead[bytes]) -> TypeGuard[IO[bytes]]: diff --git a/optuna_dashboard/artifact/exceptions.py b/optuna_dashboard/artifact/exceptions.py new file mode 100644 index 00000000..abc014df --- /dev/null +++ b/optuna_dashboard/artifact/exceptions.py @@ -0,0 +1,9 @@ +class ArtifactNotFound(Exception): + """Exception raised when an artifact is not found. + + It is typically raised while calling + :meth:`~optuna_dashboard.artifact.protocol.ArtifactBackend.open` or + :meth:`~optuna_dashboard.artifact.protocol.ArtifactBackend.remove` methods. + """ + + ... diff --git a/optuna_dashboard/artifact/file_system.py b/optuna_dashboard/artifact/file_system.py index 06139eef..7da8483c 100644 --- a/optuna_dashboard/artifact/file_system.py +++ b/optuna_dashboard/artifact/file_system.py @@ -2,11 +2,14 @@ from __future__ import annotations import os import shutil -from typing import BinaryIO from typing import TYPE_CHECKING +from optuna_dashboard.artifact.exceptions import ArtifactNotFound + if TYPE_CHECKING: + from typing import BinaryIO + from _typeshed import SupportsRead @@ -34,7 +37,11 @@ class FileSystemBackend: def open(self, artifact_id: str) -> BinaryIO: filepath = os.path.join(self._base_path, artifact_id) - return open(filepath, "rb") + try: + f = open(filepath, "rb") + except FileNotFoundError as e: + raise ArtifactNotFound("not found") from e + return f def write(self, artifact_id: str, content_body: SupportsRead[bytes]) -> None: filepath = os.path.join(self._base_path, artifact_id) @@ -43,7 +50,10 @@ class FileSystemBackend: def remove(self, artifact_id: str) -> None: filepath = os.path.join(self._base_path, artifact_id) - os.remove(filepath) + try: + os.remove(filepath) + except FileNotFoundError as e: + raise ArtifactNotFound("not found") from e if TYPE_CHECKING: diff --git a/optuna_dashboard/artifact/protocol.py b/optuna_dashboard/artifact/protocol.py index 3780395d..1355a8b2 100644 --- a/optuna_dashboard/artifact/protocol.py +++ b/optuna_dashboard/artifact/protocol.py @@ -1,20 +1,61 @@ from __future__ import annotations -from typing import BinaryIO -from typing import Protocol from typing import TYPE_CHECKING +try: + from typing import Protocol +except ImportError: + from typing_extensions import Protocol + + if TYPE_CHECKING: + from typing import BinaryIO + from _typeshed import SupportsRead class ArtifactBackend(Protocol): + """A protocol defining the interface for an artifact backend. + + An artifact backend is responsible for managing the storage and retrieval + of artifact data. The backend should provide methods for opening, writing + and removing artifacts. + """ + def open(self, artifact_id: str) -> BinaryIO: + """Open the artifact identified by the artifact_id. + + This method should return a binary file-like object in read mode, similar to + ``open(..., mode="rb")``. If the artifact does not exist, an + :exc:`~optuna_dashboard.artifact.exceptions.ArtifactNotFound` exception + should be raised. + + Args: + artifact_id: The identifier of the artifact to open. + + Returns: + BinaryIO: A binary file-like object that can be read from. + """ ... def write(self, artifact_id: str, content_body: SupportsRead[bytes]) -> None: + """Save the content to the backend. + + Args: + artifact_id: The identifier of the artifact to write to. + content_body: The content to write to the artifact. + """ ... def remove(self, artifact_id: str) -> None: + """Remove the artifact identified by the artifact_id. + + This method should delete the artifact from the backend. If the artifact does not + exist, an :exc:`~optuna_dashboard.artifact.exceptions.ArtifactNotFound` exception + may be raised. + + Args: + artifact_id: The identifier of the artifact to remove. + """ ... diff --git a/python_tests/test_boto3_artifact.py b/python_tests/test_boto3_artifact.py index 617c11e3..df5cffab 100644 --- a/python_tests/test_boto3_artifact.py +++ b/python_tests/test_boto3_artifact.py @@ -6,6 +6,7 @@ from unittest.mock import patch import boto3 from moto import mock_s3 from optuna_dashboard.artifact.boto3 import Boto3Backend +from optuna_dashboard.artifact.exceptions import ArtifactNotFound @mock_s3 @@ -63,3 +64,8 @@ class Boto3BackendTestCase(TestCase): backend.remove(artifact_id) objects = self.s3_client.list_objects(Bucket=self.bucket_name).get("Contents", []) assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 0 + + def test_file_not_found_exception(self) -> None: + backend = Boto3Backend(self.bucket_name) + with self.assertRaises(ArtifactNotFound): + backend.open("not-found-id") diff --git a/python_tests/test_file_system_artifact.py b/python_tests/test_file_system_artifact.py index 572a88cd..80cd7624 100644 --- a/python_tests/test_file_system_artifact.py +++ b/python_tests/test_file_system_artifact.py @@ -3,6 +3,7 @@ import tempfile from unittest import TestCase from optuna_dashboard.artifact.file_system import FileSystemBackend +from optuna_dashboard.artifact.exceptions import ArtifactNotFound class FileSystemBackendTestCase(TestCase): @@ -20,3 +21,10 @@ class FileSystemBackendTestCase(TestCase): with backend.open(artifact_id) as f: actual = f.read() self.assertEqual(actual, dummy_content) + + def test_file_not_found(self) -> None: + backend = FileSystemBackend(self.dir.name) + with self.assertRaises(ArtifactNotFound): + backend.open("not-found-id") + with self.assertRaises(ArtifactNotFound): + backend.remove("not-found-id")