Initial scenario tests infra (#4118)

skyvern/forge/sdk/api/real_azure.py

@@ -0,0 +1,127 @@
+"""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)