diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index ed2e2435..ed8cd549 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -526,16 +526,16 @@ class AgentDB: token: str, ) -> OrganizationAuthToken: async with self.Session() as session: - token = OrganizationAuthTokenModel( + auth_token = OrganizationAuthTokenModel( organization_id=organization_id, token_type=token_type, token=token, ) - session.add(token) + session.add(auth_token) await session.commit() - await session.refresh(token) + await session.refresh(auth_token) - return convert_to_organization_auth_token(token) + return convert_to_organization_auth_token(auth_token) async def get_artifacts_for_task_step( self, diff --git a/skyvern/forge/sdk/services/org_auth_token_service.py b/skyvern/forge/sdk/services/org_auth_token_service.py new file mode 100644 index 00000000..7f7227a6 --- /dev/null +++ b/skyvern/forge/sdk/services/org_auth_token_service.py @@ -0,0 +1,39 @@ +from datetime import timedelta + +import structlog + +from skyvern.forge.app import DATABASE +from skyvern.forge.sdk.core import security +from skyvern.forge.sdk.models import OrganizationAuthToken, OrganizationAuthTokenType + +LOG = structlog.get_logger() +API_KEY_LIFETIME = timedelta(weeks=5200) + + +async def create_org_api_token(org_id: str) -> OrganizationAuthToken: + """Creates an API token for the specified org_id. + + Args: + org_id: The org_id for which to create an API token. + + Returns: + The API token created for the specified org_id. + """ + # get the organization + organization = await DATABASE.get_organization(org_id) + if not organization: + raise Exception(f"Organization id {org_id} not found") + + # [START create_org_api_token] + api_key = security.create_access_token( + org_id, + expires_delta=API_KEY_LIFETIME, + ) + # generate OrganizationAutoToken + org_auth_token = await DATABASE.create_org_auth_token( + organization_id=org_id, + token=api_key, + token_type=OrganizationAuthTokenType.api, + ) + LOG.info(f"Created API token for organization", organization_id=org_id) + return org_auth_token