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

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
This commit is contained in:
Shuchang Zheng
2025-02-14 00:00:19 +08:00
committed by GitHub
parent b411af56a6
commit 4407c19417
11 changed files with 394 additions and 1 deletions

View File

@@ -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 ###

View File

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

View File

@@ -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:

View File

@@ -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 = (

View File

@@ -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

View File

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

View File

@@ -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

View File

@@ -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(

View File

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

View File

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

View File

@@ -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(