瀏覽代碼

Use DefaultAzureCredential for managed identity in azure blob extention (#11559)

te-chan 2 月之前
父節點
當前提交
fac83e14bc
共有 1 個文件被更改,包括 9 次插入0 次删除
  1. 9 0
      api/extensions/storage/azure_blob_storage.py

+ 9 - 0
api/extensions/storage/azure_blob_storage.py

@@ -1,6 +1,7 @@
 from collections.abc import Generator
 from datetime import UTC, datetime, timedelta
 
+from azure.identity import DefaultAzureCredential
 from azure.storage.blob import AccountSasPermissions, BlobServiceClient, ResourceTypes, generate_account_sas
 
 from configs import dify_config
@@ -18,6 +19,11 @@ class AzureBlobStorage(BaseStorage):
         self.account_name = dify_config.AZURE_BLOB_ACCOUNT_NAME
         self.account_key = dify_config.AZURE_BLOB_ACCOUNT_KEY
 
+        if self.account_key == "managedidentity":
+            self.credential = DefaultAzureCredential()
+        else:
+            self.credential = None
+
     def save(self, filename, data):
         client = self._sync_client()
         blob_container = client.get_container_client(container=self.bucket_name)
@@ -57,6 +63,9 @@ class AzureBlobStorage(BaseStorage):
         blob_container.delete_blob(filename)
 
     def _sync_client(self):
+        if self.account_key == "managedidentity":
+            return BlobServiceClient(account_url=self.account_url, credential=self.credential)
+
         cache_key = "azure_blob_sas_token_{}_{}".format(self.account_name, self.account_key)
         cache_result = redis_client.get(cache_key)
         if cache_result is not None: