浏览代码

feat: add AWS workload identity support

Sara Angel-Murphy 2 月之前
父节点
当前提交
5e873bc643
共有 2 个文件被更改,包括 37 次插入12 次删除
  1. 26 12
      backend/open_webui/storage/provider.py
  2. 11 0
      backend/open_webui/test/apps/webui/storage/test_provider.py

+ 26 - 12
backend/open_webui/storage/provider.py

@@ -101,19 +101,33 @@ class LocalStorageProvider(StorageProvider):
 
 
 class S3StorageProvider(StorageProvider):
 class S3StorageProvider(StorageProvider):
     def __init__(self):
     def __init__(self):
-        self.s3_client = boto3.client(
-            "s3",
-            region_name=S3_REGION_NAME,
-            endpoint_url=S3_ENDPOINT_URL,
-            aws_access_key_id=S3_ACCESS_KEY_ID,
-            aws_secret_access_key=S3_SECRET_ACCESS_KEY,
-            config=Config(
-                s3={
-                    "use_accelerate_endpoint": S3_USE_ACCELERATE_ENDPOINT,
-                    "addressing_style": S3_ADDRESSING_STYLE,
-                },
-            ),
+        config = Config(
+            s3={
+                "use_accelerate_endpoint": S3_USE_ACCELERATE_ENDPOINT,
+                "addressing_style": S3_ADDRESSING_STYLE,
+            },
         )
         )
+
+        # If access key and secret are provided, use them for authentication
+        if S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY:
+            self.s3_client = boto3.client(
+                "s3",
+                region_name=S3_REGION_NAME,
+                endpoint_url=S3_ENDPOINT_URL,
+                aws_access_key_id=S3_ACCESS_KEY_ID,
+                aws_secret_access_key=S3_SECRET_ACCESS_KEY,
+                config=config,
+            )
+        else:
+            # If no explicit credentials are provided, fall back to default AWS credentials
+            # This supports workload identity (IAM roles for EC2, EKS, etc.)
+            self.s3_client = boto3.client(
+                "s3",
+                region_name=S3_REGION_NAME,
+                endpoint_url=S3_ENDPOINT_URL,
+                config=config,
+            )
+
         self.bucket_name = S3_BUCKET_NAME
         self.bucket_name = S3_BUCKET_NAME
         self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
         self.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
 
 

+ 11 - 0
backend/open_webui/test/apps/webui/storage/test_provider.py

@@ -187,6 +187,17 @@ class TestS3StorageProvider:
         assert not (upload_dir / self.filename).exists()
         assert not (upload_dir / self.filename).exists()
         assert not (upload_dir / self.filename_extra).exists()
         assert not (upload_dir / self.filename_extra).exists()
 
 
+    def test_init_without_credentials(self, monkeypatch):
+        """Test that S3StorageProvider can initialize without explicit credentials."""
+        # Temporarily unset the environment variables
+        monkeypatch.setattr(provider, "S3_ACCESS_KEY_ID", None)
+        monkeypatch.setattr(provider, "S3_SECRET_ACCESS_KEY", None)
+
+        # Should not raise an exception
+        storage = provider.S3StorageProvider()
+        assert storage.s3_client is not None
+        assert storage.bucket_name == provider.S3_BUCKET_NAME
+
 
 
 class TestGCSStorageProvider:
 class TestGCSStorageProvider:
     Storage = provider.GCSStorageProvider()
     Storage = provider.GCSStorageProvider()