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

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
2025-02-14 00:00:19 +08:00
parent b411af56a6
commit 4407c19417
11 changed files with 394 additions and 1 deletions
--- a/skyvern/forge/sdk/db/client.py
+++ b/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 = (
--- a/skyvern/forge/sdk/db/id.py
+++ b/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
--- a/skyvern/forge/sdk/db/models.py
+++ b/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)