get_valid_org_auth_tokens (#531)

Skyvern-AI · Jun 29, 2024 · 6a6119b · 6a6119b
1 parent dba3d7f
commit 6a6119b
Showing 1 changed file with 25 additions and 0 deletions.
diff --git 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,