Prevent file objects unexpectedly being closed by boto after uploading to s3. Fixes #153

Author: madewithkode

tests/storage/test_storage.py

@@ -702,6 +702,22 @@ def test_fetch_file_contents_using_s3_select_raises_errors(
             )
             assert len(file_contents) == 0
 
+    def test_non_closeable_buffered_reader_does_not_close_underlying_file(self):
+        # deliberately not using a context manager to
+        # open/read the file so we can correctly observe
+        # the closing behavior.
+        tmp = tempfile.TemporaryFile(mode="w+b")
+        tmp.write(b"data")
+        tmp.seek(0)
+        try:
+            wrapped = storage.NonCloseableBufferedReader(tmp)
+            wrapped.close()
+
+            # Should have flushed, but NOT closed
+            assert not tmp.closed
+        finally:
+            tmp.close()
+
 
 class TestMemoryFileStore:
     def setup_method(self):

xocto/storage/storage.py

@@ -125,6 +125,15 @@ class ReadableBinaryFile(Protocol):
     def read(self, size: int = ...) -> bytes: ...
 
 
+class NonCloseableBufferedReader(io.BufferedReader):
+    def __init__(self, raw: ReadableBinaryFile, *args: Any, **kwargs: Any):
+        super().__init__(cast(io.RawIOBase, raw), *args, **kwargs)
+
+    def close(self) -> None:
+        # Don't close the underlying file object, just flush the buffer.
+        self.flush()
+
+
 class StreamingBodyIOAdapter(io.RawIOBase):
     """
     Wrapper to adapt a boto3 S3 object to a standard Python "file-like" object
@@ -467,6 +476,16 @@ def store_file(
 
         readable = _to_stream(contents=contents)
 
+        # If `readable` is an actual file-like object(like files from `open()` or Django’s
+        # `UploadedFile`) wrap it with `NonCloseableBufferedReader` to prevent closing
+        # by boto3 after upload. No need to do this for BytesIO/StringIO because they are
+        # in-memory buffers thus do not necessarily need protection from `close()`. Also,
+        # `BufferedReader` expects a raw binary stream, not a BytesIO or StringIO.
+        if hasattr(readable, "read") and not isinstance(
+            readable, (io.BytesIO, io.StringIO)
+        ):
+            readable = NonCloseableBufferedReader(readable)
+
         # `boto_client.upload_fileobj` is type annotated with `Fileobj: BinaryIO`. However, in
         # practice the only file-like method it needs is `read(size=...)`. This cast allows us to
         # use `upload_fileobj` with any Fileobj that implements the `ReadableBinaryFile` protocol
@@ -489,6 +508,9 @@ def store_file(
             ExtraArgs=extra_args,
         )
 
+        if isinstance(file_obj, NonCloseableBufferedReader):
+            file_obj.detach()
+
         return self.bucket_name, key_path
 
     def store_versioned_file(