Delete Azure Vault secrets async (#3716)
This commit is contained in:
committed by
GitHub
parent
ef3167cb07
commit
d6bb3de2c5
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user