Delete Azure Vault secrets async (#3716)

This commit is contained in:
Stanislav Novosad
2025-10-14 20:25:22 -06:00
committed by GitHub
parent ef3167cb07
commit d6bb3de2c5
5 changed files with 48 additions and 20 deletions

View File

@@ -1,7 +1,6 @@
import uuid
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Any, AsyncGenerator, Awaitable, Callable
from typing import Awaitable, Callable
import structlog
from fastapi import FastAPI, Response, status
@@ -51,20 +50,12 @@ def custom_openapi() -> dict:
return app.openapi_schema
@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncGenerator[None, Any]:
"""Lifespan context manager for FastAPI app startup and shutdown."""
LOG.info("Server started")
yield
LOG.info("Server shutting down")
def get_agent_app() -> FastAPI:
"""
Start the agent server.
"""
app = FastAPI(lifespan=lifespan)
app = FastAPI()
# Add CORS middleware
app.add_middleware(

View File

@@ -33,7 +33,7 @@ class AsyncAzureVaultClient:
finally:
await secret_client.close()
async def create_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str:
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)

View File

@@ -205,6 +205,7 @@ async def create_credential(
include_in_schema=False,
)
async def delete_credential(
background_tasks: BackgroundTasks,
credential_id: str = Path(
...,
description="The unique identifier of the credential to delete",
@@ -226,6 +227,9 @@ async def delete_credential(
await credential_service.delete_credential(credential)
# Schedule background cleanup if the service implements it
background_tasks.add_task(credential_service.post_delete_credential_item, credential.item_id)
return None

View File

@@ -1,6 +1,7 @@
import uuid
from typing import Annotated, Literal, Union
import structlog
from azure.identity.aio import ClientSecretCredential
from fastapi import HTTPException
from pydantic import BaseModel, Field, TypeAdapter
@@ -21,6 +22,8 @@ from skyvern.forge.sdk.schemas.credentials import (
)
from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
LOG = structlog.get_logger()
class AzureCredentialVaultService(CredentialVaultService):
class _PasswordCredentialDataImage(BaseModel):
@@ -72,7 +75,37 @@ class AzureCredentialVaultService(CredentialVaultService):
credential: Credential,
) -> None:
await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id)
await self.delete_credential_item(credential.item_id)
# Deleting takes several seconds, so we empty the value and delete async so customers do not have to wait
await self._client.create_or_update_secret(
vault_name=self._vault_name,
secret_name=credential.item_id,
secret_value="",
)
async def post_delete_credential_item(self, item_id: str) -> None:
"""
Background task to delete the credential item from Azure Key Vault.
This allows the API to respond quickly while the deletion happens asynchronously.
"""
try:
LOG.info(
"Deleting credential item from Azure Key Vault in background",
item_id=item_id,
vault_name=self._vault_name,
)
await self._client.delete_secret(secret_name=item_id, vault_name=self._vault_name)
LOG.info(
"Successfully deleted credential item from Azure Key Vault",
item_id=item_id,
vault_name=self._vault_name,
)
except Exception as e:
LOG.exception(
"Failed to delete credential item from Azure Key Vault in background",
item_id=item_id,
vault_name=self._vault_name,
error=str(e),
)
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
credential = await app.DATABASE.get_credential(credential_id=credential_id, organization_id=organization_id)
@@ -85,12 +118,6 @@ class AzureCredentialVaultService(CredentialVaultService):
credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size)
return [_convert_to_response(credential) for credential in credentials]
async def delete_credential_item(self, item_id: str) -> None:
await self._client.delete_secret(
vault_name=self._vault_name,
secret_name=item_id,
)
async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
secret_json_str = await self._client.get_secret(secret_name=db_credential.item_id, vault_name=self._vault_name)
if secret_json_str is None:
@@ -154,7 +181,7 @@ class AzureCredentialVaultService(CredentialVaultService):
secret_name = f"{organization_id}-{uuid.uuid4()}".replace("_", "")
secret_value = data.model_dump_json(exclude_none=True)
return await self._client.create_secret(
return await self._client.create_or_update_secret(
vault_name=self._vault_name,
secret_name=secret_name,
secret_value=secret_value,

View File

@@ -26,6 +26,12 @@ class CredentialVaultService(ABC):
async def delete_credential(self, credential: Credential) -> None:
"""Delete a credential from the vault and database."""
async def post_delete_credential_item(self, item_id: str) -> None:
"""
Optional hook for scheduling background cleanup tasks after credential deletion.
Default implementation does nothing. Override in subclasses as needed.
"""
@abstractmethod
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
"""Retrieve a credential with masked sensitive data."""