Skip to content

Commit 60c67b3

Browse files
authored
fix: prevent stream closure on exceptions in S3Writer (#331)
Enhanced __exit__ to log exceptions without closing stream, preventing masked errors during S3 operations. Added logging and unit tests to verify exception handling behavior.
1 parent a5f0a80 commit 60c67b3

File tree

2 files changed

+35
-1
lines changed

2 files changed

+35
-1
lines changed

s3torchconnector/src/s3torchconnector/s3writer.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,12 @@
44
import io
55
from typing import Union
66
import threading
7+
import logging
78

89
from s3torchconnectorclient._mountpoint_s3_client import PutObjectStream
910

11+
logger = logging.getLogger(__name__)
12+
1013

1114
class S3Writer(io.BufferedIOBase):
1215
"""A write-only, file like representation of a single object stored in S3."""
@@ -22,7 +25,16 @@ def __enter__(self):
2225
return self
2326

2427
def __exit__(self, exc_type, exc_val, exc_tb):
25-
self.close()
28+
"""Close stream on normal exit, log any exceptions that occurred."""
29+
if exc_type is not None:
30+
try:
31+
logger.info(
32+
f"Exception occurred before closing stream: {exc_type.__name__}: {exc_val}"
33+
)
34+
except:
35+
pass
36+
else:
37+
self.close()
2638

2739
def write(
2840
self,

s3torchconnector/tst/unit/test_s3writer.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,3 +87,25 @@ def test_concurrent_close_calls():
8787

8888
MOCK_STREAM.close.assert_called_once()
8989
assert writer._closed
90+
91+
92+
def test_exit_without_exception():
93+
"""Test __exit__ method when no exception occurs."""
94+
MOCK_STREAM.reset_mock()
95+
96+
writer = S3Writer(MOCK_STREAM)
97+
writer.__exit__(None, None, None)
98+
99+
MOCK_STREAM.close.assert_called_once()
100+
101+
102+
def test_exit_with_exception(caplog):
103+
"""Test __exit__ method when an exception occurs."""
104+
MOCK_STREAM.reset_mock()
105+
106+
writer = S3Writer(MOCK_STREAM)
107+
test_exception = ValueError("Test exception")
108+
writer.__exit__(ValueError, test_exception, None)
109+
110+
# Stream should not be closed on exception
111+
MOCK_STREAM.close.assert_not_called()

0 commit comments

Comments
 (0)