Bladeren bron

Better AWS S3 storage support (#4195)

Robert Brennan 1 jaar geleden
bovenliggende
commit
641a15356f
1 gewijzigde bestanden met toevoegingen van 25 en 8 verwijderingen
  1. 25 8
      openhands/storage/s3.py

+ 25 - 8
openhands/storage/s3.py

@@ -1,27 +1,44 @@
+import io
 import os
 
 from minio import Minio
 
 from openhands.storage.files import FileStore
 
-AWS_S3_ENDPOINT = 's3.amazonaws.com'
-
 
 class S3FileStore(FileStore):
-    def __init__(self, endpoint: str = AWS_S3_ENDPOINT) -> None:
+    def __init__(self) -> None:
         access_key = os.getenv('AWS_ACCESS_KEY_ID')
         secret_key = os.getenv('AWS_SECRET_ACCESS_KEY')
+        endpoint = os.getenv('AWS_S3_ENDPOINT', 's3.amazonaws.com')
+        secure = os.getenv('AWS_S3_SECURE', 'true').lower() == 'true'
         self.bucket = os.getenv('AWS_S3_BUCKET')
-        self.client = Minio(endpoint, access_key, secret_key)
+        self.client = Minio(endpoint, access_key, secret_key, secure=secure)
 
     def write(self, path: str, contents: str) -> None:
-        self.client.put_object(self.bucket, path, contents)
+        as_bytes = contents.encode('utf-8')
+        stream = io.BytesIO(as_bytes)
+        try:
+            self.client.put_object(self.bucket, path, stream, len(as_bytes))
+        except Exception as e:
+            raise FileNotFoundError(f'Failed to write to S3 at path {path}: {e}')
 
     def read(self, path: str) -> str:
-        return self.client.get_object(self.bucket, path).data.decode('utf-8')
+        try:
+            return self.client.get_object(self.bucket, path).data.decode('utf-8')
+        except Exception as e:
+            raise FileNotFoundError(f'Failed to read from S3 at path {path}: {e}')
 
     def list(self, path: str) -> list[str]:
-        return [obj.object_name for obj in self.client.list_objects(self.bucket, path)]
+        try:
+            return [
+                obj.object_name for obj in self.client.list_objects(self.bucket, path)
+            ]
+        except Exception as e:
+            raise FileNotFoundError(f'Failed to list S3 objects at path {path}: {e}')
 
     def delete(self, path: str) -> None:
-        self.client.remove_object(self.bucket, path)
+        try:
+            self.client.remove_object(self.bucket, path)
+        except Exception as e:
+            raise FileNotFoundError(f'Failed to delete S3 object at path {path}: {e}')