Skip to content

Commit

Permalink
fixing s3 upload close()
Browse files Browse the repository at this point in the history
  • Loading branch information
kamikaze committed Mar 6, 2023
1 parent 907b8cc commit 02af154
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions src/aiofm/protocols/s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def __init__(self, bucket_name: str, object_key: str, s3_client):
self.bucket_name = bucket_name
self.object_key = object_key
self.s3_client = s3_client
self.flushed_to_s3 = False

def write(self, data):
if isinstance(data, StreamingBody):
Expand All @@ -58,11 +59,14 @@ def write(self, data):
else:
raise ValueError(f'Unsupported data type: {type(data)}')

def _flush(self):
self.seek(0)
self.s3_client.upload_fileobj(self, Bucket=self.bucket_name, Key=self.object_key)

def close(self):
if self.tell() > 0:
self.flush()
self.seek(0)
self.s3_client.upload_fileobj(self, Bucket=self.bucket_name, Key=self.object_key)
if not self.flushed_to_s3 and self.tell() > 0:
self.flushed_to_s3 = True
self._flush()
self.seek(0)
self.truncate()
super().close()
Expand Down

0 comments on commit 02af154

Please sign in to comment.