From 5e873bc643c92394c567a9ccffdebc0035852457 Mon Sep 17 00:00:00 2001
From: Sara Angel-Murphy <sara.angelmurphy@bluevoyant.com>
Date: Thu, 27 Feb 2025 13:12:54 -0500
Subject: [PATCH] feat: add AWS workload identity support

---
 backend/open_webui/storage/provider.py        | 38 +++++++++++++------
 .../test/apps/webui/storage/test_provider.py  | 11 ++++++
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/backend/open_webui/storage/provider.py b/backend/open_webui/storage/provider.py
index 2f31cbdaf..c5c0056cc 100644
--- a/backend/open_webui/storage/provider.py
+++ b/backend/open_webui/storage/provider.py
@@ -101,19 +101,33 @@ class LocalStorageProvider(StorageProvider):
 
 class S3StorageProvider(StorageProvider):
     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.key_prefix = S3_KEY_PREFIX if S3_KEY_PREFIX else ""
 
diff --git a/backend/open_webui/test/apps/webui/storage/test_provider.py b/backend/open_webui/test/apps/webui/storage/test_provider.py
index a5ef13504..3c874592f 100644
--- a/backend/open_webui/test/apps/webui/storage/test_provider.py
+++ b/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_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:
     Storage = provider.GCSStorageProvider()