128 lines
5.5 KiB
Python
128 lines
5.5 KiB
Python
"""Real implementations of Azure clients (Vault and Storage) and their factories."""
|
|
|
|
from typing import Self
|
|
|
|
import structlog
|
|
from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
|
|
from azure.keyvault.secrets.aio import SecretClient
|
|
from azure.storage.blob.aio import BlobServiceClient
|
|
|
|
from skyvern.forge.sdk.api.azure import AsyncAzureStorageClient, AsyncAzureVaultClient, AzureClientFactory
|
|
from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential
|
|
|
|
LOG = structlog.get_logger()
|
|
|
|
|
|
class RealAsyncAzureVaultClient(AsyncAzureVaultClient):
|
|
"""Real implementation of Azure Vault client using Azure SDK."""
|
|
|
|
def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None:
|
|
self.credential = credential
|
|
|
|
async def __aenter__(self) -> Self:
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
|
|
) -> None:
|
|
await self.credential.close()
|
|
|
|
async def get_secret(self, secret_name: str, vault_name: str) -> str | None:
|
|
secret_client = await self._get_secret_client(vault_name)
|
|
try:
|
|
secret = await secret_client.get_secret(secret_name)
|
|
return secret.value
|
|
except Exception as e:
|
|
LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e)
|
|
return None
|
|
finally:
|
|
await secret_client.close()
|
|
|
|
async def create_or_update_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str:
|
|
secret_client = await self._get_secret_client(vault_name)
|
|
try:
|
|
secret = await secret_client.set_secret(secret_name, secret_value)
|
|
return secret.name
|
|
except Exception as e:
|
|
LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e)
|
|
raise e
|
|
finally:
|
|
await secret_client.close()
|
|
|
|
async def delete_secret(self, secret_name: str, vault_name: str) -> str:
|
|
secret_client = await self._get_secret_client(vault_name)
|
|
try:
|
|
secret = await secret_client.delete_secret(secret_name)
|
|
return secret.name
|
|
except Exception as e:
|
|
LOG.exception("Failed to delete secret from Azure Key Vault.", secret_name=secret_name, error=e)
|
|
raise e
|
|
finally:
|
|
await secret_client.close()
|
|
|
|
async def _get_secret_client(self, vault_name: str) -> SecretClient:
|
|
# 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_name}.vault.azure.net" # Placeholder, adjust as needed
|
|
return SecretClient(vault_url=key_vault_url, credential=self.credential)
|
|
|
|
async def close(self) -> None:
|
|
await self.credential.close()
|
|
|
|
|
|
class RealAsyncAzureStorageClient(AsyncAzureStorageClient):
|
|
"""Real implementation of Azure Storage client using Azure SDK."""
|
|
|
|
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
|
|
try:
|
|
await container_client.create_container()
|
|
except Exception as e:
|
|
LOG.info("Azure container already exists or failed to create", container_name=container_name, error=e)
|
|
|
|
with open(file_path, "rb") as data:
|
|
await container_client.upload_blob(name=blob_name, data=data, overwrite=True)
|
|
LOG.info("File uploaded to Azure Blob Storage", container_name=container_name, blob_name=blob_name)
|
|
except Exception as e:
|
|
LOG.error(
|
|
"Failed to upload file to Azure Blob Storage",
|
|
container_name=container_name,
|
|
blob_name=blob_name,
|
|
error=e,
|
|
)
|
|
raise e
|
|
|
|
async def close(self) -> None:
|
|
await self.blob_service_client.close()
|
|
|
|
|
|
class RealAzureClientFactory(AzureClientFactory):
|
|
"""Factory for creating real Azure Vault and Storage clients."""
|
|
|
|
def create_default(self) -> AsyncAzureVaultClient:
|
|
"""Create an Azure Vault client using DefaultAzureCredential."""
|
|
return RealAsyncAzureVaultClient(DefaultAzureCredential())
|
|
|
|
def create_from_client_secret(self, credential: AzureClientSecretCredential) -> AsyncAzureVaultClient:
|
|
"""Create an Azure Vault client using client secret credentials."""
|
|
cred = ClientSecretCredential(
|
|
tenant_id=credential.tenant_id,
|
|
client_id=credential.client_id,
|
|
client_secret=credential.client_secret,
|
|
)
|
|
return RealAsyncAzureVaultClient(cred)
|
|
|
|
def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> AsyncAzureStorageClient:
|
|
"""Create an Azure Storage client with the provided credentials."""
|
|
return RealAsyncAzureStorageClient(storage_account_name, storage_account_key)
|