Add credentials table, CRUD endpoints, and credential parameter (#1767)

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>

alembic/versions/2025_02_13_1554-26c5ed737819_added_credentials_table_and_credential_.py

@@ -0,0 +1,70 @@
+"""Added credentials table and credential parameter
+
+Revision ID: 26c5ed737819
+Revises: b111f0f795bd
+Create Date: 2025-02-13 15:54:32.388064+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "26c5ed737819"
+down_revision: Union[str, None] = "b111f0f795bd"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "credentials",
+        sa.Column("credential_id", sa.String(), nullable=False),
+        sa.Column("organization_id", sa.String(), nullable=False),
+        sa.Column("credential_type", sa.String(), nullable=False),
+        sa.Column("name", sa.String(), nullable=False),
+        sa.Column("website_url", sa.String(), nullable=True),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("modified_at", sa.DateTime(), nullable=False),
+        sa.Column("deleted_at", sa.DateTime(), nullable=True),
+        sa.PrimaryKeyConstraint("credential_id"),
+    )
+    op.create_table(
+        "credential_parameters",
+        sa.Column("credential_parameter_id", sa.String(), nullable=False),
+        sa.Column("workflow_id", sa.String(), nullable=False),
+        sa.Column("key", sa.String(), nullable=False),
+        sa.Column("description", sa.String(), nullable=True),
+        sa.Column("credential_id", sa.String(), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("modified_at", sa.DateTime(), nullable=False),
+        sa.Column("deleted_at", sa.DateTime(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["workflow_id"],
+            ["workflows.workflow_id"],
+        ),
+        sa.PrimaryKeyConstraint("credential_parameter_id"),
+    )
+    op.create_index(
+        op.f("ix_credential_parameters_credential_parameter_id"),
+        "credential_parameters",
+        ["credential_parameter_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_credential_parameters_workflow_id"), "credential_parameters", ["workflow_id"], unique=False
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_credential_parameters_workflow_id"), table_name="credential_parameters")
+    op.drop_index(op.f("ix_credential_parameters_credential_parameter_id"), table_name="credential_parameters")
+    op.drop_table("credential_parameters")
+    op.drop_table("credentials")
+    # ### end Alembic commands ###

skyvern/exceptions.py

@@ -309,6 +309,16 @@ class BitwardenAccessDeniedError(BitwardenBaseError):
         )
 
 
+class CredentialParameterParsingError(SkyvernException):
+    def __init__(self, message: str) -> None:
+        super().__init__(f"Error parsing credential parameter: {message}")
+
+
+class CredentialParameterNotFoundError(SkyvernException):
+    def __init__(self, credential_parameter_id: str) -> None:
+        super().__init__(f"Could not find credential parameter: {credential_parameter_id}")
+
+
 class UnknownElementTreeFormat(SkyvernException):
     def __init__(self, fmt: str) -> None:
         super().__init__(f"Unknown element tree format {fmt}")

skyvern/forge/sdk/api/aws.py

@@ -44,6 +44,30 @@ class AsyncAWSClient:
             LOG.exception("Failed to get secret.", secret_name=secret_name, error_code=error_code)
             return None
 
+    @execute_with_async_client(client_type=AWSClientType.SECRETS_MANAGER)
+    async def create_secret(self, secret_name: str, secret_value: str, client: AioBaseClient = None) -> None:
+        try:
+            await client.create_secret(Name=secret_name, SecretString=secret_value)
+        except Exception as e:
+            LOG.exception("Failed to create secret.", secret_name=secret_name)
+            raise e
+
+    @execute_with_async_client(client_type=AWSClientType.SECRETS_MANAGER)
+    async def set_secret(self, secret_name: str, secret_value: str, client: AioBaseClient = None) -> None:
+        try:
+            await client.put_secret_value(SecretId=secret_name, SecretString=secret_value)
+        except Exception as e:
+            LOG.exception("Failed to set secret.", secret_name=secret_name)
+            raise e
+
+    @execute_with_async_client(client_type=AWSClientType.SECRETS_MANAGER)
+    async def delete_secret(self, secret_name: str, client: AioBaseClient = None) -> None:
+        try:
+            await client.delete_secret(SecretId=secret_name)
+        except Exception as e:
+            LOG.exception("Failed to delete secret.", secret_name=secret_name)
+            raise e
+
     @execute_with_async_client(client_type=AWSClientType.S3)
     async def upload_file(self, uri: str, data: bytes, client: AioBaseClient = None) -> str | None:
         try:

skyvern/forge/sdk/db/client.py

@@ -20,6 +20,8 @@ from skyvern.forge.sdk.db.models import (
     BitwardenCreditCardDataParameterModel,
     BitwardenLoginCredentialParameterModel,
     BitwardenSensitiveInformationParameterModel,
+    CredentialModel,
+    CredentialParameterModel,
     ObserverCruiseModel,
     ObserverThoughtModel,
     OrganizationAuthTokenModel,
@@ -59,6 +61,7 @@ from skyvern.forge.sdk.db.utils import (
 from skyvern.forge.sdk.log_artifacts import save_workflow_run_logs
 from skyvern.forge.sdk.models import Step, StepStatus
 from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
+from skyvern.forge.sdk.schemas.credentials import Credential, CredentialType
 from skyvern.forge.sdk.schemas.observers import ObserverTask, ObserverTaskStatus, ObserverThought, ObserverThoughtType
 from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
@@ -73,6 +76,7 @@ from skyvern.forge.sdk.workflow.models.parameter import (
     BitwardenCreditCardDataParameter,
     BitwardenLoginCredentialParameter,
     BitwardenSensitiveInformationParameter,
+    CredentialParameter,
     OutputParameter,
     WorkflowParameter,
     WorkflowParameterType,
@@ -1666,6 +1670,30 @@ class AgentDB:
             LOG.error("SQLAlchemyError", exc_info=True)
             raise
 
+    async def create_credential_parameter(
+        self, workflow_id: str, key: str, credential_id: str, description: str | None = None
+    ) -> CredentialParameter:
+        async with self.Session() as session:
+            credential_parameter = CredentialParameterModel(
+                workflow_id=workflow_id,
+                key=key,
+                description=description,
+                credential_id=credential_id,
+            )
+            session.add(credential_parameter)
+            await session.commit()
+            await session.refresh(credential_parameter)
+            return CredentialParameter(
+                credential_parameter_id=credential_parameter.credential_parameter_id,
+                workflow_id=credential_parameter.workflow_id,
+                key=credential_parameter.key,
+                description=credential_parameter.description,
+                credential_id=credential_parameter.credential_id,
+                created_at=credential_parameter.created_at,
+                modified_at=credential_parameter.modified_at,
+                deleted_at=credential_parameter.deleted_at,
+            )
+
     async def get_workflow_run_output_parameters(self, workflow_run_id: str) -> list[WorkflowRunOutputParameter]:
         try:
             async with self.Session() as session:
@@ -2673,6 +2701,84 @@ class AgentDB:
             await session.refresh(task_run)
             return TaskRun.model_validate(task_run)
 
+    async def create_credential(
+        self, name: str, website_url: str | None, credential_type: CredentialType, organization_id: str
+    ) -> Credential:
+        async with self.Session() as session:
+            credential = CredentialModel(
+                organization_id=organization_id,
+                name=name,
+                website_url=website_url,
+                credential_type=credential_type,
+            )
+            session.add(credential)
+            await session.commit()
+            await session.refresh(credential)
+            return Credential.model_validate(credential)
+
+    async def get_credential(self, credential_id: str, organization_id: str) -> Credential:
+        async with self.Session() as session:
+            credential = (
+                await session.scalars(
+                    select(CredentialModel)
+                    .filter_by(credential_id=credential_id)
+                    .filter_by(organization_id=organization_id)
+                    .filter(CredentialModel.deleted_at.is_(None))
+                )
+            ).first()
+            if credential:
+                return Credential.model_validate(credential)
+            raise NotFoundError(f"Credential {credential_id} not found")
+
+    async def get_credentials(self, organization_id: str) -> list[Credential]:
+        async with self.Session() as session:
+            credentials = (
+                await session.scalars(
+                    select(CredentialModel)
+                    .filter_by(organization_id=organization_id)
+                    .filter(CredentialModel.deleted_at.is_(None))
+                    .order_by(CredentialModel.created_at.desc())
+                )
+            ).all()
+            return [Credential.model_validate(credential) for credential in credentials]
+
+    async def update_credential(
+        self, credential_id: str, organization_id: str, name: str | None = None, website_url: str | None = None
+    ) -> Credential:
+        async with self.Session() as session:
+            credential = (
+                await session.scalars(
+                    select(CredentialModel)
+                    .filter_by(credential_id=credential_id)
+                    .filter_by(organization_id=organization_id)
+                )
+            ).first()
+            if not credential:
+                raise NotFoundError(f"Credential {credential_id} not found")
+            if name:
+                credential.name = name
+            if website_url:
+                credential.website_url = website_url
+            await session.commit()
+            await session.refresh(credential)
+            return Credential.model_validate(credential)
+
+    async def delete_credential(self, credential_id: str, organization_id: str) -> None:
+        async with self.Session() as session:
+            credential = (
+                await session.scalars(
+                    select(CredentialModel)
+                    .filter_by(credential_id=credential_id)
+                    .filter_by(organization_id=organization_id)
+                )
+            ).first()
+            if not credential:
+                raise NotFoundError(f"Credential {credential_id} not found")
+            credential.deleted_at = datetime.utcnow()
+            await session.commit()
+            await session.refresh(credential)
+            return None
+
     async def cache_task_run(self, run_id: str, organization_id: str | None = None) -> TaskRun:
         async with self.Session() as session:
             task_run = (

skyvern/forge/sdk/db/id.py

@@ -34,6 +34,7 @@ AWS_SECRET_PARAMETER_PREFIX = "asp"
 BITWARDEN_CREDIT_CARD_DATA_PARAMETER_PREFIX = "bccd"
 BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
 BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi"
+CREDENTIAL_PARAMETER_PREFIX = "cp"
 OBSERVER_CRUISE_ID = "oc"
 OBSERVER_THOUGHT_ID = "ot"
 ORGANIZATION_AUTH_TOKEN_PREFIX = "oat"
@@ -51,6 +52,7 @@ WORKFLOW_PERMANENT_ID_PREFIX = "wpid"
 WORKFLOW_PREFIX = "w"
 WORKFLOW_RUN_BLOCK_PREFIX = "wrb"
 WORKFLOW_RUN_PREFIX = "wr"
+CREDENTIAL_PREFIX = "cred"
 
 
 def generate_workflow_id() -> str:
@@ -173,6 +175,16 @@ def generate_task_run_id() -> str:
     return f"{TASK_RUN_PREFIX}_{int_id}"
 
 
+def generate_credential_id() -> str:
+    int_id = generate_id()
+    return f"{CREDENTIAL_PREFIX}_{int_id}"
+
+
+def generate_credential_parameter_id() -> str:
+    int_id = generate_id()
+    return f"{CREDENTIAL_PARAMETER_PREFIX}_{int_id}"
+
+
 def generate_id() -> int:
     """
     generate a 64-bit int ID

skyvern/forge/sdk/db/models.py

@@ -27,6 +27,8 @@ from skyvern.forge.sdk.db.id import (
     generate_bitwarden_credit_card_data_parameter_id,
     generate_bitwarden_login_credential_parameter_id,
     generate_bitwarden_sensitive_information_parameter_id,
+    generate_credential_id,
+    generate_credential_parameter_id,
     generate_observer_cruise_id,
     generate_observer_thought_id,
     generate_org_id,
@@ -382,6 +384,21 @@ class BitwardenCreditCardDataParameterModel(Base):
     deleted_at = Column(DateTime, nullable=True)
 
 
+class CredentialParameterModel(Base):
+    __tablename__ = "credential_parameters"
+
+    credential_parameter_id = Column(String, primary_key=True, index=True, default=generate_credential_parameter_id)
+    workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False)
+    key = Column(String, nullable=False)
+    description = Column(String, nullable=True)
+
+    credential_id = Column(String, nullable=False)
+
+    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    deleted_at = Column(DateTime, nullable=True)
+
+
 class WorkflowRunParameterModel(Base):
     __tablename__ = "workflow_run_parameters"
 
@@ -629,3 +646,18 @@ class TaskRunModel(Base):
     cached = Column(Boolean, nullable=False, default=False)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
     modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+
+
+class CredentialModel(Base):
+    __tablename__ = "credentials"
+
+    credential_id = Column(String, primary_key=True, default=generate_credential_id)
+    organization_id = Column(String, nullable=False)
+
+    credential_type = Column(String, nullable=False)
+    name = Column(String, nullable=False)
+    website_url = Column(String, nullable=True)
+
+    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    deleted_at = Column(DateTime, nullable=True)

skyvern/forge/sdk/schemas/credentials.py

@@ -0,0 +1,49 @@
+from datetime import datetime
+from enum import StrEnum
+
+from pydantic import BaseModel, ConfigDict
+
+
+class CredentialType(StrEnum):
+    PASSWORD = "password"
+    CREDIT_CARD = "credit_card"
+
+
+class PasswordCredential(BaseModel):
+    password: str
+    username: str
+
+
+class CreditCardCredential(BaseModel):
+    card_number: str
+    card_cvv: str
+    card_exp_month: str
+    card_exp_year: str
+    card_brand: str
+    card_holder_name: str
+
+
+class UpdateCredentialRequest(BaseModel):
+    name: str | None = None
+    website_url: str | None = None
+
+
+class CreateCredentialRequest(BaseModel):
+    name: str
+    website_url: str | None = None
+    credential_type: CredentialType
+    credential: PasswordCredential | CreditCardCredential
+
+
+class Credential(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
+    credential_id: str
+    organization_id: str
+    name: str
+    website_url: str | None = None
+    credential_type: CredentialType
+
+    created_at: datetime
+    modified_at: datetime
+    deleted_at: datetime | None = None

skyvern/forge/sdk/workflow/context_manager.py

@@ -1,10 +1,17 @@
+import json
 import uuid
 from typing import TYPE_CHECKING, Any, Self
 
 import structlog
 
 from skyvern.config import settings
-from skyvern.exceptions import BitwardenBaseError, SkyvernException, WorkflowRunContextNotInitialized
+from skyvern.exceptions import (
+    BitwardenBaseError,
+    CredentialParameterNotFoundError,
+    CredentialParameterParsingError,
+    SkyvernException,
+    WorkflowRunContextNotInitialized,
+)
 from skyvern.forge.sdk.api.aws import AsyncAWSClient
 from skyvern.forge.sdk.schemas.organizations import Organization
 from skyvern.forge.sdk.schemas.tasks import TaskStatus
@@ -17,6 +24,7 @@ from skyvern.forge.sdk.workflow.models.parameter import (
     BitwardenLoginCredentialParameter,
     BitwardenSensitiveInformationParameter,
     ContextParameter,
+    CredentialParameter,
     OutputParameter,
     Parameter,
     ParameterType,
@@ -45,6 +53,7 @@ class WorkflowRunContext:
             | BitwardenLoginCredentialParameter
             | BitwardenCreditCardDataParameter
             | BitwardenSensitiveInformationParameter
+            | CredentialParameter
         ],
     ) -> Self:
         # key is label name
@@ -68,6 +77,10 @@ class WorkflowRunContext:
         for secrete_parameter in secret_parameters:
             if isinstance(secrete_parameter, AWSSecretParameter):
                 await workflow_run_context.register_aws_secret_parameter_value(aws_client, secrete_parameter)
+            elif isinstance(secrete_parameter, CredentialParameter):
+                await workflow_run_context.register_credential_parameter_value(
+                    aws_client, secrete_parameter, organization
+                )
             elif isinstance(secrete_parameter, BitwardenLoginCredentialParameter):
                 await workflow_run_context.register_bitwarden_login_credential_parameter_value(
                     aws_client, secrete_parameter, organization
@@ -161,6 +174,37 @@ class WorkflowRunContext:
     def generate_random_secret_id() -> str:
         return f"secret_{uuid.uuid4()}"
 
+    async def register_credential_parameter_value(
+        self,
+        aws_client: AsyncAWSClient,
+        parameter: CredentialParameter,
+        organization: Organization,
+    ) -> None:
+        LOG.info(f"Fetching credential parameter value for credential: {parameter.credential_id}")
+        org_secret_values = await aws_client.get_secret(organization.organization_id)
+        if org_secret_values is None:
+            raise CredentialParameterNotFoundError(parameter.credential_id)
+        # Parse the items and extract credentials
+        try:
+            org_secret_values_json = json.loads(org_secret_values)
+
+        except json.JSONDecodeError:
+            raise CredentialParameterParsingError(
+                f"Failed to parse credential JSON. Credential ID: {parameter.credential_id}"
+            )
+
+        credentials = org_secret_values_json.get(parameter.credential_id)
+        if credentials is None:
+            raise CredentialParameterNotFoundError(parameter.credential_id)
+
+        self.parameters[parameter.key] = parameter
+        self.values[parameter.key] = {}
+        for key, value in credentials.items():
+            random_secret_id = self.generate_random_secret_id()
+            secret_id = f"{random_secret_id}_{key}"
+            self.secrets[secret_id] = value
+            self.values[parameter.key][key] = secret_id
+
     async def register_aws_secret_parameter_value(
         self,
         aws_client: AsyncAWSClient,
@@ -550,6 +594,7 @@ class WorkflowRunContext:
                     BitwardenLoginCredentialParameter,
                     BitwardenCreditCardDataParameter,
                     BitwardenSensitiveInformationParameter,
+                    CredentialParameter,
                 ),
             ):
                 LOG.error(

skyvern/forge/sdk/workflow/models/parameter.py

@@ -17,6 +17,7 @@ class ParameterType(StrEnum):
     BITWARDEN_SENSITIVE_INFORMATION = "bitwarden_sensitive_information"
     BITWARDEN_CREDIT_CARD_DATA = "bitwarden_credit_card_data"
     OUTPUT = "output"
+    CREDENTIAL = "credential"
 
 
 class Parameter(BaseModel, abc.ABC):
@@ -65,6 +66,20 @@ class BitwardenLoginCredentialParameter(Parameter):
     deleted_at: datetime | None = None
 
 
+class CredentialParameter(Parameter):
+    model_config = ConfigDict(from_attributes=True)
+    parameter_type: Literal[ParameterType.CREDENTIAL] = ParameterType.CREDENTIAL
+
+    credential_parameter_id: str
+    workflow_id: str
+
+    credential_id: str
+
+    created_at: datetime
+    modified_at: datetime
+    deleted_at: datetime | None = None
+
+
 class BitwardenSensitiveInformationParameter(Parameter):
     parameter_type: Literal[ParameterType.BITWARDEN_SENSITIVE_INFORMATION] = (
         ParameterType.BITWARDEN_SENSITIVE_INFORMATION
@@ -182,5 +197,6 @@ ParameterSubclasses = Union[
     BitwardenSensitiveInformationParameter,
     BitwardenCreditCardDataParameter,
     OutputParameter,
+    CredentialParameter,
 ]
 PARAMETER_TYPE = Annotated[ParameterSubclasses, Field(discriminator="parameter_type")]

skyvern/forge/sdk/workflow/models/yaml.py

@@ -43,6 +43,11 @@ class BitwardenLoginCredentialParameterYAML(ParameterYAML):
     bitwarden_collection_id: str | None = None
 
 
+class CredentialParameterYAML(ParameterYAML):
+    parameter_type: Literal[ParameterType.CREDENTIAL] = ParameterType.CREDENTIAL  # type: ignore
+    credential_id: str
+
+
 class BitwardenSensitiveInformationParameterYAML(ParameterYAML):
     # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
     # Parameter 1 of Literal[...] cannot be of type "Any"
@@ -341,6 +346,7 @@ PARAMETER_YAML_SUBCLASSES = (
     | WorkflowParameterYAML
     | ContextParameterYAML
     | OutputParameterYAML
+    | CredentialParameterYAML
 )
 PARAMETER_YAML_TYPES = Annotated[PARAMETER_YAML_SUBCLASSES, Field(discriminator="parameter_type")]
 

skyvern/forge/sdk/workflow/service.py

@@ -64,6 +64,7 @@ from skyvern.forge.sdk.workflow.models.parameter import (
     BitwardenLoginCredentialParameter,
     BitwardenSensitiveInformationParameter,
     ContextParameter,
+    CredentialParameter,
     OutputParameter,
     Parameter,
     ParameterType,
@@ -233,6 +234,7 @@ class WorkflowService:
                     BitwardenLoginCredentialParameter,
                     BitwardenCreditCardDataParameter,
                     BitwardenSensitiveInformationParameter,
+                    CredentialParameter,
                 ),
             )
         ]
@@ -817,6 +819,20 @@ class WorkflowService:
             bitwarden_collection_id=bitwarden_collection_id,
         )
 
+    async def create_credential_parameter(
+        self,
+        workflow_id: str,
+        key: str,
+        credential_id: str,
+        description: str | None = None,
+    ) -> CredentialParameter:
+        return await app.DATABASE.create_credential_parameter(
+            workflow_id=workflow_id,
+            key=key,
+            credential_id=credential_id,
+            description=description,
+        )
+
     async def create_bitwarden_sensitive_information_parameter(
         self,
         workflow_id: str,
@@ -1358,6 +1374,13 @@ class WorkflowService:
                         key=parameter.key,
                         description=parameter.description,
                     )
+                elif parameter.parameter_type == ParameterType.CREDENTIAL:
+                    parameters[parameter.key] = await self.create_credential_parameter(
+                        workflow_id=workflow.workflow_id,
+                        key=parameter.key,
+                        description=parameter.description,
+                        credential_id=parameter.credential_id,
+                    )
                 elif parameter.parameter_type == ParameterType.BITWARDEN_LOGIN_CREDENTIAL:
                     if not parameter.bitwarden_collection_id:
                         raise WorkflowParameterMissingRequiredValue(