From d80d49131cf185237678005211014965945e596e Mon Sep 17 00:00:00 2001 From: Stanislav Novosad Date: Tue, 4 Nov 2025 10:29:51 -0700 Subject: [PATCH] Clean up Azure migration (#3895) Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> --- skyvern/forge/sdk/routes/credentials.py | 49 ++++++++--- .../azure_credential_vault_service.py | 42 --------- .../bitwarden_credential_service.py | 87 ------------------- .../credential/credential_vault_service.py | 9 -- 4 files changed, 39 insertions(+), 148 deletions(-) diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py index bf0302f6..7302e194 100644 --- a/skyvern/forge/sdk/routes/credentials.py +++ b/skyvern/forge/sdk/routes/credentials.py @@ -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") diff --git a/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py b/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py index 949d1457..629d8df3 100644 --- a/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py +++ b/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py @@ -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") diff --git a/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py b/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py index 78caefbb..af930d16 100644 --- a/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py +++ b/skyvern/forge/sdk/services/credential/bitwarden_credential_service.py @@ -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) diff --git a/skyvern/forge/sdk/services/credential/credential_vault_service.py b/skyvern/forge/sdk/services/credential/credential_vault_service.py index 93c16b19..3e70134d 100644 --- a/skyvern/forge/sdk/services/credential/credential_vault_service.py +++ b/skyvern/forge/sdk/services/credential/credential_vault_service.py @@ -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."""