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 import uuid
from contextlib import asynccontextmanager
from datetime import datetime from datetime import datetime
from typing import Any, AsyncGenerator, Awaitable, Callable from typing import Awaitable, Callable
import structlog import structlog
from fastapi import FastAPI, Response, status from fastapi import FastAPI, Response, status
@@ -51,20 +50,12 @@ def custom_openapi() -> dict:
return app.openapi_schema 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: def get_agent_app() -> FastAPI:
""" """
Start the agent server. Start the agent server.
""" """
app = FastAPI(lifespan=lifespan) app = FastAPI()
# Add CORS middleware # Add CORS middleware
app.add_middleware( app.add_middleware(

View File

@@ -33,7 +33,7 @@ class AsyncAzureVaultClient:
finally: finally:
await secret_client.close() 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) secret_client = await self._get_secret_client(vault_name)
try: try:
secret = await secret_client.set_secret(secret_name, secret_value) secret = await secret_client.set_secret(secret_name, secret_value)

View File

@@ -205,6 +205,7 @@ async def create_credential(
include_in_schema=False, include_in_schema=False,
) )
async def delete_credential( async def delete_credential(
background_tasks: BackgroundTasks,
credential_id: str = Path( credential_id: str = Path(
..., ...,
description="The unique identifier of the credential to delete", description="The unique identifier of the credential to delete",
@@ -226,6 +227,9 @@ async def delete_credential(
await credential_service.delete_credential(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 return None

View File

@@ -1,6 +1,7 @@
import uuid import uuid
from typing import Annotated, Literal, Union from typing import Annotated, Literal, Union
import structlog
from azure.identity.aio import ClientSecretCredential from azure.identity.aio import ClientSecretCredential
from fastapi import HTTPException from fastapi import HTTPException
from pydantic import BaseModel, Field, TypeAdapter 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 from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
LOG = structlog.get_logger()
class AzureCredentialVaultService(CredentialVaultService): class AzureCredentialVaultService(CredentialVaultService):
class _PasswordCredentialDataImage(BaseModel): class _PasswordCredentialDataImage(BaseModel):
@@ -72,7 +75,37 @@ class AzureCredentialVaultService(CredentialVaultService):
credential: Credential, credential: Credential,
) -> None: ) -> None:
await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id) 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: 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) 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) credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size)
return [_convert_to_response(credential) for credential in credentials] 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: 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) 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: if secret_json_str is None:
@@ -154,7 +181,7 @@ class AzureCredentialVaultService(CredentialVaultService):
secret_name = f"{organization_id}-{uuid.uuid4()}".replace("_", "") secret_name = f"{organization_id}-{uuid.uuid4()}".replace("_", "")
secret_value = data.model_dump_json(exclude_none=True) 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, vault_name=self._vault_name,
secret_name=secret_name, secret_name=secret_name,
secret_value=secret_value, secret_value=secret_value,

View File

@@ -26,6 +26,12 @@ class CredentialVaultService(ABC):
async def delete_credential(self, credential: Credential) -> None: async def delete_credential(self, credential: Credential) -> None:
"""Delete a credential from the vault and database.""" """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 @abstractmethod
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse: async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
"""Retrieve a credential with masked sensitive data.""" """Retrieve a credential with masked sensitive data."""