Clean up Azure migration (#3895)

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
This commit is contained in:
Stanislav Novosad
2025-11-04 10:29:51 -07:00
committed by GitHub
parent b1fe444984
commit d80d49131c
4 changed files with 39 additions and 148 deletions

View File

@@ -21,6 +21,7 @@ from skyvern.forge.sdk.routes.code_samples import (
from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
from skyvern.forge.sdk.schemas.credentials import (
CreateCredentialRequest,
Credential,
CredentialResponse,
CredentialType,
CredentialVaultType,
@@ -226,7 +227,7 @@ async def create_credential(
),
current_org: Organization = Depends(org_auth_service.get_current_org),
) -> CredentialResponse:
credential_service = await _get_credential_vault_service(current_org.organization_id)
credential_service = await _get_credential_vault_service()
credential = await credential_service.create_credential(organization_id=current_org.organization_id, data=data)
@@ -348,9 +349,13 @@ async def get_credential(
),
current_org: Organization = Depends(org_auth_service.get_current_org),
) -> CredentialResponse:
credential_service = await _get_credential_vault_service(current_org.organization_id)
credential = await app.DATABASE.get_credential(
credential_id=credential_id, organization_id=current_org.organization_id
)
if not credential:
raise HTTPException(status_code=404, detail="Credential not found")
return await credential_service.get_credential(current_org.organization_id, credential_id)
return _convert_to_response(credential)
@legacy_base_router.get("/credentials")
@@ -395,9 +400,8 @@ async def get_credentials(
openapi_extra={"x-fern-sdk-parameter-name": "page_size"},
),
) -> list[CredentialResponse]:
credential_service = await _get_credential_vault_service(current_org.organization_id)
return await credential_service.get_credentials(current_org.organization_id, page, page_size)
credentials = await app.DATABASE.get_credentials(current_org.organization_id, page=page, page_size=page_size)
return [_convert_to_response(credential) for credential in credentials]
@base_router.get(
@@ -606,10 +610,8 @@ async def update_azure_client_secret_credential(
)
async def _get_credential_vault_service(organization_id: str) -> CredentialVaultService:
org_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
if settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.BITWARDEN or org_collection:
async def _get_credential_vault_service() -> CredentialVaultService:
if settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.BITWARDEN:
return app.BITWARDEN_CREDENTIAL_VAULT_SERVICE
elif settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.AZURE_VAULT:
if not app.AZURE_CREDENTIAL_VAULT_SERVICE:
@@ -617,3 +619,30 @@ async def _get_credential_vault_service(organization_id: str) -> CredentialVault
return app.AZURE_CREDENTIAL_VAULT_SERVICE
else:
raise HTTPException(status_code=400, detail="Credential storage not supported")
def _convert_to_response(credential: Credential) -> CredentialResponse:
if credential.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(
username=credential.username or credential.credential_id,
totp_type=credential.totp_type,
)
return CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=credential.credential_type,
name=credential.name,
)
elif credential.credential_type == CredentialType.CREDIT_CARD:
credential_response = CreditCardCredentialResponse(
last_four=credential.card_last4 or "****",
brand=credential.card_brand or "Card Brand",
)
return CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=credential.credential_type,
name=credential.name,
)
else:
raise HTTPException(status_code=400, detail="Credential type not supported")

View File

@@ -3,7 +3,6 @@ from typing import Annotated, Literal, Union
import structlog
from azure.identity.aio import ClientSecretCredential
from fastapi import HTTPException
from pydantic import BaseModel, Field, TypeAdapter
from skyvern.forge import app
@@ -12,13 +11,10 @@ from skyvern.forge.sdk.schemas.credentials import (
CreateCredentialRequest,
Credential,
CredentialItem,
CredentialResponse,
CredentialType,
CredentialVaultType,
CreditCardCredential,
CreditCardCredentialResponse,
PasswordCredential,
PasswordCredentialResponse,
)
from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
@@ -107,17 +103,6 @@ class AzureCredentialVaultService(CredentialVaultService):
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)
if not credential:
raise HTTPException(status_code=404, detail="Credential not found")
return _convert_to_response(credential)
async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
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 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:
@@ -186,30 +171,3 @@ class AzureCredentialVaultService(CredentialVaultService):
secret_name=secret_name,
secret_value=secret_value,
)
def _convert_to_response(credential: Credential) -> CredentialResponse:
if credential.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(
username=credential.username or credential.credential_id,
totp_type=credential.totp_type,
)
return CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=credential.credential_type,
name=credential.name,
)
elif credential.credential_type == CredentialType.CREDIT_CARD:
credential_response = CreditCardCredentialResponse(
last_four=credential.card_last4 or "****",
brand=credential.card_brand or "Card Brand",
)
return CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=credential.credential_type,
name=credential.name,
)
else:
raise HTTPException(status_code=400, detail="Credential type not supported")

View File

@@ -6,11 +6,7 @@ from skyvern.forge.sdk.schemas.credentials import (
CreateCredentialRequest,
Credential,
CredentialItem,
CredentialResponse,
CredentialType,
CredentialVaultType,
CreditCardCredentialResponse,
PasswordCredentialResponse,
)
from skyvern.forge.sdk.services.bitwarden import BitwardenService
from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
@@ -63,88 +59,5 @@ class BitwardenCredentialVaultService(CredentialVaultService):
await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id)
await BitwardenService.delete_credential_item(credential.item_id)
async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
if not organization_bitwarden_collection:
raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
credential = await app.DATABASE.get_credential(credential_id=credential_id, organization_id=organization_id)
if not credential:
raise HTTPException(status_code=404, detail="Credential not found")
credential_item = await BitwardenService.get_credential_item(credential.item_id)
if not credential_item:
raise HTTPException(status_code=404, detail="Credential not found")
if credential_item.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(
username=credential_item.credential.username,
totp_type=credential.totp_type,
)
return CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=credential_item.credential_type,
name=credential_item.name,
)
if credential_item.credential_type == CredentialType.CREDIT_CARD:
credential_response = CreditCardCredentialResponse(
last_four=credential_item.credential.card_number[-4:],
brand=credential_item.credential.card_brand,
)
return CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=credential_item.credential_type,
name=credential_item.name,
)
raise HTTPException(status_code=400, detail="Invalid credential type")
async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
if not organization_bitwarden_collection:
return []
credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size)
items = await BitwardenService.get_collection_items(organization_bitwarden_collection.collection_id)
response_items = []
for credential in credentials:
item = next((item for item in items if item.item_id == credential.item_id), None)
if not item:
LOG.warning(
"Credential item not found in vault",
credential_id=credential.credential_id,
item_id=credential.item_id,
)
continue
if item.credential_type == CredentialType.PASSWORD:
credential_response = PasswordCredentialResponse(
username=item.credential.username,
totp_type=credential.totp_type,
)
response_items.append(
CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=item.credential_type,
name=item.name,
)
)
elif item.credential_type == CredentialType.CREDIT_CARD:
credential_response = CreditCardCredentialResponse(
last_four=item.credential.card_number[-4:],
brand=item.credential.card_brand,
)
response_items.append(
CredentialResponse(
credential=credential_response,
credential_id=credential.credential_id,
credential_type=item.credential_type,
name=item.name,
)
)
return response_items
async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
return await BitwardenService.get_credential_item(db_credential.item_id)

View File

@@ -5,7 +5,6 @@ from skyvern.forge.sdk.schemas.credentials import (
CreateCredentialRequest,
Credential,
CredentialItem,
CredentialResponse,
CredentialType,
CredentialVaultType,
)
@@ -32,14 +31,6 @@ class CredentialVaultService(ABC):
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."""
@abstractmethod
async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
"""Retrieve all credentials for an organization with pagination."""
@abstractmethod
async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
"""Retrieve the full credential data from the vault."""