Azure ClientSecretCredential support (#3456)

Co-authored-by: Suchintan <suchintan@users.noreply.github.com>
Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>
This commit is contained in:
stenn930
2025-09-23 10:16:48 -06:00
committed by GitHub
parent 10fac9bad0
commit a29a2bc49b
12 changed files with 592 additions and 71 deletions

View File

@@ -1,39 +1,24 @@
import structlog
from azure.identity.aio import DefaultAzureCredential
from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
from azure.keyvault.secrets.aio import SecretClient
from azure.storage.blob.aio import BlobServiceClient
from skyvern.exceptions import AzureConfigurationError
from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential
LOG = structlog.get_logger()
class AsyncAzureClient:
def __init__(self, storage_account_name: str | None, storage_account_key: str | None):
self.storage_account_name = storage_account_name
self.storage_account_key = storage_account_key
if storage_account_name and storage_account_key:
self.blob_service_client = BlobServiceClient(
account_url=f"https://{storage_account_name}.blob.core.windows.net",
credential=storage_account_key,
)
else:
self.blob_service_client = None
self.credential = DefaultAzureCredential()
async def get_secret(self, secret_name: str, vault_name: str | None = None) -> str | None:
vault_subdomain = vault_name or self.storage_account_name
if not vault_subdomain:
raise AzureConfigurationError("Missing vault")
class AsyncAzureVaultClient:
def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential):
self.credential = credential
async def get_secret(self, secret_name: str, vault_name: str) -> str | None:
try:
# Azure Key Vault URL format: https://<your-key-vault-name>.vault.azure.net
# Assuming the secret_name is actually the Key Vault URL and the secret name
# This needs to be clarified or passed as separate parameters
# For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings.
key_vault_url = f"https://{vault_subdomain}.vault.azure.net" # Placeholder, adjust as needed
key_vault_url = f"https://{vault_name}.vault.azure.net" # Placeholder, adjust as needed
secret_client = SecretClient(vault_url=key_vault_url, credential=self.credential)
secret = await secret_client.get_secret(secret_name)
return secret.value
@@ -43,10 +28,34 @@ class AsyncAzureClient:
finally:
await self.credential.close()
async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None:
if not self.blob_service_client:
raise AzureConfigurationError("Storage is not configured")
async def close(self) -> None:
await self.credential.close()
@classmethod
def create_default(cls) -> "AsyncAzureVaultClient":
return cls(DefaultAzureCredential())
@classmethod
def create_from_client_secret(
cls,
credential: AzureClientSecretCredential,
) -> "AsyncAzureVaultClient":
cred = ClientSecretCredential(
tenant_id=credential.tenant_id,
client_id=credential.client_id,
client_secret=credential.client_secret,
)
return cls(cred)
class AsyncAzureStorageClient:
def __init__(self, storage_account_name: str, storage_account_key: str):
self.blob_service_client = BlobServiceClient(
account_url=f"https://{storage_account_name}.blob.core.windows.net",
credential=storage_account_key,
)
async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None:
try:
container_client = self.blob_service_client.get_container_client(container_name)
# Create the container if it doesn't exist
@@ -69,4 +78,3 @@ class AsyncAzureClient:
async def close(self) -> None:
await self.blob_service_client.close()
await self.credential.close()