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

This commit is contained in:
te-chan 2025-02-05 14:43:43 +09:00 committed by GitHub
parent a97cec57e4
commit fac83e14bc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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: