Skip to content

Commit

Permalink
Use DefaultAzureCredential for managed identity in azure blob extenti…
Browse files Browse the repository at this point in the history
  • Loading branch information
te-chan2 authored Feb 5, 2025
1 parent a97cec5 commit fac83e1
Showing 1 changed file with 9 additions and 0 deletions.
9 changes: 9 additions & 0 deletions api/extensions/storage/azure_blob_storage.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -18,6 +19,11 @@ def __init__(self):
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)
Expand Down Expand Up @@ -57,6 +63,9 @@ def delete(self, filename):
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:
Expand Down

0 comments on commit fac83e1

Please sign in to comment.