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

@@ -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)
+    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]
-    return await credential_service.get_credentials(current_org.organization_id, page, page_size)
 
 
 @base_router.get(
@@ -606,10 +610,8 @@ async def update_azure_client_secret_credential(
         )
 
 
-async def _get_credential_vault_service(organization_id: str) -> CredentialVaultService:
+async def _get_credential_vault_service() -> CredentialVaultService:
-    org_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
+    if settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.BITWARDEN:
-
-    if settings.CREDENTIAL_VAULT_TYPE == CredentialVaultType.BITWARDEN or org_collection:
         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")

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")

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)

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."""