Initial scenario tests infra (#4118)
This commit is contained in:
Stanislav Novosad
committed by
GitHub
GitHub
parent
4dcc8ed725
commit
5b530cab52
@@ -1,115 +1,126 @@
|
||||
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 typing import Protocol, Self
|
||||
|
||||
from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential
|
||||
|
||||
LOG = structlog.get_logger()
|
||||
|
||||
class AsyncAzureVaultClient(Protocol):
|
||||
"""Protocol defining the interface for Azure Vault clients.
|
||||
|
||||
class AsyncAzureVaultClient:
|
||||
def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None:
|
||||
self.credential = credential
|
||||
This client provides methods to interact with Azure Key Vault for secret management.
|
||||
"""
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
return self
|
||||
"""Enter async context manager."""
|
||||
...
|
||||
|
||||
async def __aexit__(
|
||||
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
|
||||
) -> None:
|
||||
await self.credential.close()
|
||||
"""Exit async context manager and cleanup resources."""
|
||||
...
|
||||
|
||||
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()
|
||||
"""Retrieve a secret from Azure Key Vault.
|
||||
|
||||
Args:
|
||||
secret_name: The name of the secret to retrieve
|
||||
vault_name: The name of the Azure Key Vault
|
||||
|
||||
Returns:
|
||||
The secret value as a string, or None if the secret doesn't exist or an error occurs
|
||||
"""
|
||||
...
|
||||
|
||||
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()
|
||||
"""Create or update a secret in Azure Key Vault.
|
||||
|
||||
Args:
|
||||
secret_name: The name of the secret to create or update
|
||||
secret_value: The value to store
|
||||
vault_name: The name of the Azure Key Vault
|
||||
|
||||
Returns:
|
||||
The name of the created/updated secret
|
||||
|
||||
Raises:
|
||||
Exception: If the operation fails
|
||||
"""
|
||||
...
|
||||
|
||||
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()
|
||||
"""Delete a secret from Azure Key Vault.
|
||||
|
||||
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)
|
||||
Args:
|
||||
secret_name: The name of the secret to delete
|
||||
vault_name: The name of the Azure Key Vault
|
||||
|
||||
Returns:
|
||||
The name of the deleted secret
|
||||
|
||||
Raises:
|
||||
Exception: If the operation fails
|
||||
"""
|
||||
...
|
||||
|
||||
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)
|
||||
"""Close the client and release all resources."""
|
||||
...
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
class AsyncAzureStorageClient(Protocol):
|
||||
"""Protocol defining the interface for Azure Storage clients.
|
||||
|
||||
This client provides methods to interact with Azure Blob Storage for file operations.
|
||||
"""
|
||||
|
||||
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)
|
||||
"""Upload a file from the local filesystem to Azure Blob Storage.
|
||||
|
||||
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
|
||||
Args:
|
||||
container_name: The name of the Azure Blob container
|
||||
blob_name: The name to give the blob in storage
|
||||
file_path: The local path to the file to upload
|
||||
|
||||
Raises:
|
||||
Exception: If the upload fails
|
||||
"""
|
||||
...
|
||||
|
||||
async def close(self) -> None:
|
||||
await self.blob_service_client.close()
|
||||
"""Close the storage client and release resources."""
|
||||
...
|
||||
|
||||
|
||||
class AzureClientFactory(Protocol):
|
||||
"""Protocol defining the interface for creating Azure Vault and Storage clients."""
|
||||
|
||||
def create_default(self) -> "AsyncAzureVaultClient":
|
||||
"""Create an Azure Vault client using default credentials.
|
||||
|
||||
Returns:
|
||||
An AsyncAzureVaultClient instance using DefaultAzureCredential
|
||||
"""
|
||||
...
|
||||
|
||||
def create_from_client_secret(self, credential: AzureClientSecretCredential) -> "AsyncAzureVaultClient":
|
||||
"""Create an Azure Vault client using client secret credentials.
|
||||
|
||||
Args:
|
||||
credential: Azure client secret credentials containing tenant_id, client_id, and client_secret
|
||||
|
||||
Returns:
|
||||
An AsyncAzureVaultClient instance
|
||||
"""
|
||||
...
|
||||
|
||||
def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> "AsyncAzureStorageClient":
|
||||
"""Create an Azure Storage client with the provided credentials.
|
||||
|
||||
Args:
|
||||
storage_account_name: The name of the Azure storage account
|
||||
storage_account_key: The access key for the storage account
|
||||
|
||||
Returns:
|
||||
An AsyncAzureStorageClient instance
|
||||
"""
|
||||
...
|
||||
|
||||
Reference in New Issue
Block a user