Azure Vault credential support (#3394)

skyvern/forge/sdk/db/client.py

@@ -17,6 +17,7 @@ from skyvern.forge.sdk.db.models import (
     AISuggestionModel,
     ArtifactModel,
     AWSSecretParameterModel,
+    AzureVaultCredentialParameterModel,
     BitwardenCreditCardDataParameterModel,
     BitwardenLoginCredentialParameterModel,
     BitwardenSensitiveInformationParameterModel,
@@ -88,6 +89,7 @@ from skyvern.forge.sdk.schemas.totp_codes import TOTPCode
 from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
 from skyvern.forge.sdk.workflow.models.parameter import (
     AWSSecretParameter,
+    AzureVaultCredentialParameter,
     BitwardenCreditCardDataParameter,
     BitwardenLoginCredentialParameter,
     BitwardenSensitiveInformationParameter,
@@ -2107,6 +2109,43 @@ class AgentDB:
                 deleted_at=parameter.deleted_at,
             )
 
+    async def create_azure_vault_credential_parameter(
+        self,
+        workflow_id: str,
+        key: str,
+        vault_name: str,
+        username_key: str,
+        password_key: str,
+        totp_secret_key: str | None = None,
+        description: str | None = None,
+    ) -> AzureVaultCredentialParameter:
+        async with self.Session() as session:
+            parameter = AzureVaultCredentialParameterModel(
+                workflow_id=workflow_id,
+                key=key,
+                description=description,
+                vault_name=vault_name,
+                username_key=username_key,
+                password_key=password_key,
+                totp_secret_key=totp_secret_key,
+            )
+            session.add(parameter)
+            await session.commit()
+            await session.refresh(parameter)
+            return AzureVaultCredentialParameter(
+                azure_vault_credential_parameter_id=parameter.azure_vault_credential_parameter_id,
+                workflow_id=parameter.workflow_id,
+                key=parameter.key,
+                description=parameter.description,
+                vault_name=parameter.vault_name,
+                username_key=parameter.username_key,
+                password_key=parameter.password_key,
+                totp_secret_key=parameter.totp_secret_key,
+                created_at=parameter.created_at,
+                modified_at=parameter.modified_at,
+                deleted_at=parameter.deleted_at,
+            )
+
     async def get_workflow_run_output_parameters(self, workflow_run_id: str) -> list[WorkflowRunOutputParameter]:
         try:
             async with self.Session() as session:

skyvern/forge/sdk/db/id.py

@@ -35,6 +35,7 @@ BITWARDEN_CREDIT_CARD_DATA_PARAMETER_PREFIX = "bccd"
 BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
 BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi"
 CREDENTIAL_ONEPASSWORD_PARAMETER_PREFIX = "opp"
+CREDENTIAL_AZURE_VAULT_PARAMETER_PREFIX = "azcp"
 CREDENTIAL_PARAMETER_PREFIX = "cp"
 CREDENTIAL_PREFIX = "cred"
 DEBUG_SESSION_PREFIX = "ds"
@@ -123,6 +124,11 @@ def generate_onepassword_credential_parameter_id() -> str:
     return f"{CREDENTIAL_ONEPASSWORD_PARAMETER_PREFIX}_{int_id}"
 
 
+def generate_azure_vault_credential_parameter_id() -> str:
+    int_id = generate_id()
+    return f"{CREDENTIAL_AZURE_VAULT_PARAMETER_PREFIX}_{int_id}"
+
+
 def generate_organization_auth_token_id() -> str:
     int_id = generate_id()
     return f"{ORGANIZATION_AUTH_TOKEN_PREFIX}_{int_id}"

skyvern/forge/sdk/db/models.py

@@ -24,6 +24,7 @@ from skyvern.forge.sdk.db.id import (
     generate_ai_suggestion_id,
     generate_artifact_id,
     generate_aws_secret_parameter_id,
+    generate_azure_vault_credential_parameter_id,
     generate_bitwarden_credit_card_data_parameter_id,
     generate_bitwarden_login_credential_parameter_id,
     generate_bitwarden_sensitive_information_parameter_id,
@@ -467,6 +468,30 @@ class OnePasswordCredentialParameterModel(Base):
     deleted_at = Column(DateTime, nullable=True)
 
 
+class AzureVaultCredentialParameterModel(Base):
+    __tablename__ = "azure_vault_credential_parameters"
+
+    azure_vault_credential_parameter_id = Column(
+        String, primary_key=True, default=generate_azure_vault_credential_parameter_id
+    )
+    workflow_id = Column(String, index=True, nullable=False)
+    key = Column(String, nullable=False)
+    description = Column(String, nullable=True)
+    vault_name = Column(String, nullable=False)
+    username_key = Column(String, nullable=False)
+    password_key = Column(String, nullable=False)
+    totp_secret_key = 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)
+
+
 class WorkflowRunParameterModel(Base):
     __tablename__ = "workflow_run_parameters"