From ca5bf5f315e6bb0bf11e883617dd182f9349c0ba Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Tue, 13 Jun 2023 03:30:02 +0900
Subject: [PATCH] Fix a bug that boto3.upload_fileobj may close the file

---
 optuna_dashboard/artifact/boto3.py  | 18 ++++++++++++++++--
 python_tests/artifact/test_boto3.py |  4 +++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/optuna_dashboard/artifact/boto3.py b/optuna_dashboard/artifact/boto3.py
index 43abe6c0..fcddf5d7 100644
--- a/optuna_dashboard/artifact/boto3.py
+++ b/optuna_dashboard/artifact/boto3.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import io
+import shutil
 from typing import TYPE_CHECKING
 
 import boto3
@@ -33,9 +35,15 @@ class Boto3Backend:
               return ...
     """
 
-    def __init__(self, bucket_name: str, client: Optional[S3Client] = None) -> None:
+    def __init__(
+        self, bucket_name: str, client: Optional[S3Client] = None, *, avoid_buf_copy: bool = False
+    ) -> None:
         self.bucket = bucket_name
         self.client = client or boto3.client("s3")
+        # This flag is added to avoid that upload_fileobj() method of Boto3 client
+        # may close the source file object.
+        # See https://github.com/boto/boto3/issues/929
+        self._avoid_buf_copy = avoid_buf_copy
 
     def open(self, artifact_id: str) -> BinaryIO:
         try:
@@ -49,7 +57,13 @@ class Boto3Backend:
         return body  # type: ignore
 
     def write(self, artifact_id: str, content_body: BinaryIO) -> None:
-        self.client.upload_fileobj(content_body, self.bucket, artifact_id)
+        fsrc: BinaryIO = content_body
+        if not self._avoid_buf_copy:
+            buf = io.BytesIO()
+            shutil.copyfileobj(content_body, buf)
+            buf.seek(0)
+            fsrc = buf
+        self.client.upload_fileobj(fsrc, self.bucket, artifact_id)
 
     def remove(self, artifact_id: str) -> None:
         try:
diff --git a/python_tests/artifact/test_boto3.py b/python_tests/artifact/test_boto3.py
index 570bdc66..189a5bb2 100644
--- a/python_tests/artifact/test_boto3.py
+++ b/python_tests/artifact/test_boto3.py
@@ -26,9 +26,10 @@ class Boto3BackendTestCase(TestCase):
     def test_upload_download(self) -> None:
         artifact_id = "dummy-uuid"
         dummy_content = b"Hello World"
+        buf = io.BytesIO(dummy_content)
 
         backend = Boto3Backend(self.bucket_name)
-        backend.write(artifact_id, io.BytesIO(dummy_content))
+        backend.write(artifact_id, buf)
         assert len(self.s3_client.list_objects(Bucket=self.bucket_name)["Contents"]) == 1
         obj = self.s3_client.get_object(Bucket=self.bucket_name, Key=artifact_id)
         assert obj["Body"].read() == dummy_content
@@ -36,6 +37,7 @@ class Boto3BackendTestCase(TestCase):
         with backend.open(artifact_id) as f:
             actual = f.read()
         self.assertEqual(actual, dummy_content)
+        self.assertFalse(buf.closed)
 
     def test_remove(self) -> None:
         artifact_id = "dummy-uuid"