Use bitwarden server to manage credentials (#1806)

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
This commit is contained in:
Shuchang Zheng
2025-02-20 13:50:41 -08:00
committed by GitHub
parent 902c0ad4ce
commit 02a8861d4a
11 changed files with 609 additions and 35 deletions

View File

@@ -25,6 +25,7 @@ from skyvern.forge.sdk.db.models import (
ObserverCruiseModel,
ObserverThoughtModel,
OrganizationAuthTokenModel,
OrganizationBitwardenCollectionModel,
OrganizationModel,
OutputParameterModel,
PersistentBrowserSessionModel,
@@ -63,6 +64,7 @@ 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.organization_bitwarden_collections import OrganizationBitwardenCollection
from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken
from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
from skyvern.forge.sdk.schemas.task_generations import TaskGeneration
@@ -2745,14 +2747,18 @@ class AgentDB:
return TaskRun.model_validate(task_run)
async def create_credential(
self, name: str, website_url: str | None, credential_type: CredentialType, organization_id: str
self,
name: str,
credential_type: CredentialType,
organization_id: str,
item_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,
item_id=item_id,
)
session.add(credential)
await session.commit()
@@ -2773,7 +2779,7 @@ class AgentDB:
return Credential.model_validate(credential)
raise NotFoundError(f"Credential {credential_id} not found")
async def get_credentials(self, organization_id: str) -> list[Credential]:
async def get_credentials(self, organization_id: str, page: int = 1, page_size: int = 10) -> list[Credential]:
async with self.Session() as session:
credentials = (
await session.scalars(
@@ -2781,6 +2787,8 @@ class AgentDB:
.filter_by(organization_id=organization_id)
.filter(CredentialModel.deleted_at.is_(None))
.order_by(CredentialModel.created_at.desc())
.offset((page - 1) * page_size)
.limit(page_size)
)
).all()
return [Credential.model_validate(credential) for credential in credentials]
@@ -2822,6 +2830,34 @@ class AgentDB:
await session.refresh(credential)
return None
async def create_organization_bitwarden_collection(
self,
organization_id: str,
collection_id: str,
) -> OrganizationBitwardenCollection:
async with self.Session() as session:
organization_bitwarden_collection = OrganizationBitwardenCollectionModel(
organization_id=organization_id, collection_id=collection_id
)
session.add(organization_bitwarden_collection)
await session.commit()
await session.refresh(organization_bitwarden_collection)
return OrganizationBitwardenCollection.model_validate(organization_bitwarden_collection)
async def get_organization_bitwarden_collection(
self,
organization_id: str,
) -> OrganizationBitwardenCollection | None:
async with self.Session() as session:
organization_bitwarden_collection = (
await session.scalars(
select(OrganizationBitwardenCollectionModel).filter_by(organization_id=organization_id)
)
).first()
if organization_bitwarden_collection:
return OrganizationBitwardenCollection.model_validate(organization_bitwarden_collection)
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

@@ -35,6 +35,8 @@ BITWARDEN_CREDIT_CARD_DATA_PARAMETER_PREFIX = "bccd"
BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi"
CREDENTIAL_PARAMETER_PREFIX = "cp"
CREDENTIAL_PREFIX = "cred"
ORGANIZATION_BITWARDEN_COLLECTION_PREFIX = "obc"
OBSERVER_CRUISE_ID = "oc"
OBSERVER_THOUGHT_ID = "ot"
ORGANIZATION_AUTH_TOKEN_PREFIX = "oat"
@@ -52,7 +54,6 @@ 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:
@@ -175,14 +176,19 @@ def generate_task_run_id() -> str:
return f"{TASK_RUN_PREFIX}_{int_id}"
def generate_credential_parameter_id() -> str:
int_id = generate_id()
return f"{CREDENTIAL_PARAMETER_PREFIX}_{int_id}"
def generate_credential_id() -> str:
int_id = generate_id()
return f"{CREDENTIAL_PREFIX}_{int_id}"
def generate_credential_parameter_id() -> str:
def generate_organization_bitwarden_collection_id() -> str:
int_id = generate_id()
return f"{CREDENTIAL_PARAMETER_PREFIX}_{int_id}"
return f"{ORGANIZATION_BITWARDEN_COLLECTION_PREFIX}_{int_id}"
def generate_id() -> int:

View File

@@ -33,6 +33,7 @@ from skyvern.forge.sdk.db.id import (
generate_observer_thought_id,
generate_org_id,
generate_organization_auth_token_id,
generate_organization_bitwarden_collection_id,
generate_output_parameter_id,
generate_persistent_browser_session_id,
generate_step_id,
@@ -648,15 +649,29 @@ class TaskRunModel(Base):
modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
class OrganizationBitwardenCollectionModel(Base):
__tablename__ = "organization_bitwarden_collections"
organization_bitwarden_collection_id = Column(
String, primary_key=True, default=generate_organization_bitwarden_collection_id
)
organization_id = Column(String, nullable=False, index=True)
collection_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)
class CredentialModel(Base):
__tablename__ = "credentials"
credential_id = Column(String, primary_key=True, default=generate_credential_id)
organization_id = Column(String, nullable=False)
item_id = Column(String, nullable=True)
credential_type = Column(String, nullable=False)
name = Column(String, nullable=False)
website_url = Column(String, nullable=True)
credential_type = 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)