From 6a6119b406b9c46adf51d763e2c64fea620b3e2f Mon Sep 17 00:00:00 2001 From: Kerem Yilmaz Date: Fri, 28 Jun 2024 18:10:13 -0700 Subject: [PATCH] get_valid_org_auth_tokens (#531) --- skyvern/forge/sdk/db/client.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 70140a794..b18f01195 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -524,6 +524,7 @@ async def get_valid_org_auth_token( .filter_by(organization_id=organization_id) .filter_by(token_type=token_type) .filter_by(valid=True) + .order_by(OrganizationAuthTokenModel.created_at.desc()) ) ).first(): return convert_to_organization_auth_token(token) @@ -536,6 +537,30 @@ async def get_valid_org_auth_token( LOG.error("UnexpectedError", exc_info=True) raise + async def get_valid_org_auth_tokens( + self, + organization_id: str, + token_type: OrganizationAuthTokenType, + ) -> list[OrganizationAuthToken]: + try: + async with self.Session() as session: + tokens = ( + await session.scalars( + select(OrganizationAuthTokenModel) + .filter_by(organization_id=organization_id) + .filter_by(token_type=token_type) + .filter_by(valid=True) + .order_by(OrganizationAuthTokenModel.created_at.desc()) + ) + ).all() + return [convert_to_organization_auth_token(token) for token in tokens] + except SQLAlchemyError: + LOG.error("SQLAlchemyError", exc_info=True) + raise + except Exception: + LOG.error("UnexpectedError", exc_info=True) + raise + async def validate_org_auth_token( self, organization_id: str,