mage-ai 0.8.97__py3-none-any.whl → 0.8.99__py3-none-any.whl

mage_ai/api/resources/OauthResource.py

@@ -1,16 +1,22 @@
+import json
+import urllib.parse
 from datetime import datetime, timedelta
+
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.authentication.oauth2 import generate_access_token
 from mage_ai.authentication.oauth.constants import (
+    ACTIVE_DIRECTORY_CLIENT_ID,
     GITHUB_CLIENT_ID,
     GITHUB_STATE,
+    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
     OAUTH_PROVIDER_GITHUB,
+    VALID_OAUTH_PROVIDERS,
 )
 from mage_ai.authentication.oauth.utils import access_tokens_for_provider
-from mage_ai.authentication.oauth2 import generate_access_token
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
-import urllib.parse
+from mage_ai.settings import ACTIVE_DIRECTORY_DIRECTORY_ID
 
 
 class OauthResource(GenericResource):
@@ -22,7 +28,7 @@ class OauthResource(GenericResource):
         provider = payload.get('provider')
         token = payload.get('token')
 
-        if not provider or provider not in [OAUTH_PROVIDER_GITHUB]:
+        if not provider or provider not in VALID_OAUTH_PROVIDERS:
             error.update(dict(message='Invalid provider.'))
             raise ApiError(error)
 
@@ -65,7 +71,7 @@ class OauthResource(GenericResource):
         model = dict(provider=pk)
 
         error = ApiError.RESOURCE_INVALID.copy()
-        if pk not in [OAUTH_PROVIDER_GITHUB]:
+        if pk not in VALID_OAUTH_PROVIDERS:
             error.update(dict(message='Invalid provider.'))
             raise ApiError(error)
 
@@ -94,5 +100,26 @@ class OauthResource(GenericResource):
                     query_strings.append(f'{k}={v}')
 
                 model['url'] = f"https://github.com/login/oauth/authorize?{'&'.join(query_strings)}"
+            elif OAUTH_PROVIDER_ACTIVE_DIRECTORY == pk:
+                ad_directory_id = ACTIVE_DIRECTORY_DIRECTORY_ID
+                if ad_directory_id:
+                    mage_redirect_uri = '?'.join([
+                        redirect_uri,
+                        f'provider={pk}',
+                    ])
+                    query = dict(
+                        client_id=ACTIVE_DIRECTORY_CLIENT_ID,
+                        redirect_uri=f'https://api.mage.ai/v1/oauth/{pk}',
+                        response_type='code',
+                        scope='User.Read',
+                        state=urllib.parse.quote_plus(json.dumps(dict(
+                            redirect_uri=mage_redirect_uri,
+                            tenant_id=ad_directory_id,
+                        ))),
+                    )
+                    query_strings = []
+                    for k, v in query.items():
+                        query_strings.append(f'{k}={v}')
+                    model['url'] = f"https://login.microsoftonline.com/{ad_directory_id}/oauth2/v2.0/authorize?{'&'.join(query_strings)}"  # noqa: E501
 
         return self(model, user, **kwargs)

mage_ai/api/resources/PipelineResource.py

@@ -1,5 +1,6 @@
 import asyncio
 
+from sqlalchemy import or_
 from sqlalchemy.orm import aliased
 
 from mage_ai.api.operations.constants import DELETE
@@ -79,7 +80,13 @@ class PipelineResource(BaseResource):
                     a.status,
                     a.updated_at,
                 ]).
-                filter(a.pipeline_uuid.in_(pipeline_uuids))
+                filter(
+                    a.pipeline_uuid.in_(pipeline_uuids),
+                    or_(
+                        a.repo_path == get_repo_path(),
+                        a.repo_path.is_(None),
+                    )
+                )
             ).all()
             return group_by(lambda x: x.pipeline_uuid, result)
 

mage_ai/api/resources/PullRequestResource.py

@@ -0,0 +1,87 @@
+from github import Auth, Github
+from mage_ai.api.errors import ApiError
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.data_preparation.git import api
+from typing import Dict
+
+
+def pull_request_to_dict(pr) -> Dict:
+    return dict(
+        body=pr.body,
+        created_at=pr.created_at,
+        id=pr.id,
+        is_merged=pr.is_merged(),
+        last_modified=pr.last_modified,
+        merged=pr.merged,
+        state=pr.state,
+        title=pr.title,
+        url=pr.html_url,
+        user=pr.user.login,
+    )
+
+
+class PullRequestResource(GenericResource):
+    @classmethod
+    def collection(self, query, meta, user, **kwargs):
+        arr = []
+
+        repository = query.get('repository', None)
+        if repository:
+            repository = repository[0]
+
+            access_token = api.get_access_token_for_user(user)
+            if access_token:
+                auth = Auth.Token(access_token.token)
+                g = Github(auth=auth)
+                repo = g.get_repo(repository)
+                pulls = repo.get_pulls(
+                    direction='desc',
+                    sort='created',
+                    state='open',
+                ).get_page(0)
+
+                for pr in pulls:
+                    arr.append(pull_request_to_dict(pr))
+
+        return self.build_result_set(arr, user, **kwargs)
+
+    @classmethod
+    def create(self, payload, user, **kwargs):
+        error = ApiError.RESOURCE_INVALID
+
+        for key in [
+            'base_branch',
+            'compare_branch',
+            'title',
+        ]:
+            if key not in payload:
+                error.update(dict(message=f'Value for {key} is required but empty.'))
+                raise ApiError(error)
+
+        repository = payload.get('repository')
+        if not repository:
+            error.update(dict(
+                message='Repository is empty, ' +
+                'please select a repository to create a pull request in.',
+            ))
+            raise ApiError(error)
+
+        access_token = api.get_access_token_for_user(user)
+        if not access_token:
+            error.update(dict(
+                message='Access token not found, please authenticate with GitHub.',
+            ))
+            raise ApiError(error)
+
+        auth = Auth.Token(access_token.token)
+        g = Github(auth=auth)
+        repo = g.get_repo(repository)
+
+        pr = repo.create_pull(
+            base=payload.get('base_branch'),
+            body=payload.get('body'),
+            head=payload.get('compare_branch'),
+            title=payload.get('title'),
+        )
+
+        return self(pull_request_to_dict(pr), user, **kwargs)

mage_ai/api/resources/RoleResource.py

@@ -1,5 +1,5 @@
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Role
 
@@ -41,14 +41,17 @@ class RoleResource(GenericResource):
                 roles.append(permission.role)
         else:
             roles = Role.query.all()
-        access = user.get_access(Permission.Entity.PROJECT, get_repo_path())
+
+        access = 0
+        if user:
+            access = user.get_access(Permission.Entity.PROJECT, get_project_uuid())
 
         if (access & Permission.Access.OWNER == 0) and limit_roles:
             role_access = Permission.Access.EDITOR | Permission.Access.VIEWER
             roles = list(filter(
                 lambda role: role.get_access(
                     Permission.Entity.PROJECT,
-                    get_repo_path(),
+                    get_project_uuid(),
                 ) | role_access == role_access,  # Only editors and viewers
                 roles,
             ))

mage_ai/api/resources/SecretResource.py

@@ -2,15 +2,12 @@ from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
 from mage_ai.data_preparation.preferences import get_preferences
 from mage_ai.data_preparation.repo_manager import get_repo_path
-from mage_ai.data_preparation.shared.secrets import (
-    create_secret,
-    get_valid_secrets,
-)
+from mage_ai.data_preparation.shared.secrets import create_secret, get_valid_secrets
 from mage_ai.data_preparation.sync import (
-    GitConfig,
     GIT_ACCESS_TOKEN_SECRET_NAME,
     GIT_SSH_PRIVATE_KEY_SECRET_NAME,
     GIT_SSH_PUBLIC_KEY_SECRET_NAME,
+    GitConfig,
 )
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.secrets import Secret

mage_ai/api/resources/SessionResource.py

@@ -4,6 +4,8 @@ from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.BaseResource import BaseResource
 from mage_ai.authentication.ldap import new_ldap_connection
 from mage_ai.authentication.oauth2 import encode_token, generate_access_token
+from mage_ai.authentication.oauth.active_directory import get_user_info
+from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_ACTIVE_DIRECTORY
 from mage_ai.authentication.passwords import verify_password
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Role, User
@@ -18,6 +20,22 @@ class SessionResource(BaseResource):
         email = payload.get('email')
         password = payload.get('password')
         username = payload.get('username')
+        token = payload.get('token')
+        provider = payload.get('provider')
+
+        if token and provider:
+            if provider == OAUTH_PROVIDER_ACTIVE_DIRECTORY:
+                user_info = get_user_info(token)
+                principal_name = user_info.get('userPrincipalName')
+                user = User.query.filter(User.email == principal_name).first()
+                if not user:  # noqa: E712
+                    print('first user login, creating user.')
+                    user = User.create(
+                        username=principal_name,
+                        email=principal_name,
+                    )
+                oauth_token = generate_access_token(user, kwargs['oauth_client'])
+                return self(oauth_token, user, **kwargs)
 
         error = ApiError.RESOURCE_NOT_FOUND
         error.update({'message': 'Email/username and/or password invalid.'})

mage_ai/api/resources/StatusResource.py

@@ -3,8 +3,8 @@ import os
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.data_preparation.models.constants import MAX_PRINT_OUTPUT_LINES
 from mage_ai.data_preparation.repo_manager import (
-    ProjectType,
     get_project_type,
+    get_project_uuid,
     get_repo_config,
     get_repo_path,
 )
@@ -32,8 +32,10 @@ class StatusResource(GenericResource):
             KUBE_NAMESPACE,
         )
         instance_type = None
-        if get_project_type() == ProjectType.MAIN:
-            instance_type = get_repo_config().cluster_type
+        project_type = get_project_type()
+        repo_config = get_repo_config()
+        if repo_config.cluster_type:
+            instance_type = repo_config.cluster_type
         elif os.getenv(ECS_CLUSTER_NAME):
             instance_type = ClusterType.ECS
         elif os.getenv(GCP_PROJECT_ID):
@@ -56,5 +58,7 @@ class StatusResource(GenericResource):
             'disable_pipeline_edit_access': is_disable_pipeline_edit_access(),
             'max_print_output_lines': MAX_PRINT_OUTPUT_LINES,
             'require_user_authentication': REQUIRE_USER_AUTHENTICATION,
+            'project_type': project_type,
+            'project_uuid': get_project_uuid(),
         }
         return self.build_result_set([status], user, **kwargs)

mage_ai/api/resources/UserResource.py

@@ -7,7 +7,7 @@ from mage_ai.authentication.passwords import (
     generate_salt,
     verify_password,
 )
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Permission, Role, User
 from mage_ai.shared.hash import extract, ignore_keys
@@ -30,7 +30,7 @@ class UserResource(DatabaseResource):
             order_by(User.username.asc())
         )
 
-        if user.is_admin:
+        if user and user.is_admin:
             results = list(filter(lambda user: user.project_access & 3 == 0, results))
 
         return results
@@ -48,12 +48,9 @@ class UserResource(DatabaseResource):
         role_ids = payload.get('roles_new', [])
         roles_new = self.check_roles(role_ids)
 
-        missing_values = []
-        if len(roles_new) == 0:
-            missing_values.append('roles')
-
         payload['roles_new'] = roles_new
 
+        missing_values = []
         for key in ['email', 'password']:
             if not payload.get(key):
                 missing_values.append(key)
@@ -123,18 +120,13 @@ class UserResource(DatabaseResource):
             role_ids = payload.get('roles_new', [])
             roles_new = self.check_roles(role_ids)
 
-            missing_values = []
-            if len(roles_new) == 0:
-                missing_values.append('roles')
-
             payload['roles_new'] = roles_new
 
-            if len(missing_values) >= 1:
-                error.update(
-                    {'message': 'Missing required values: {}.'.format(', '.join(missing_values))})
-                raise ApiError(error)
-
-            access = get_access_for_roles(roles_new, Permission.Entity.PROJECT, get_repo_path())
+            access = get_access_for_roles(
+                roles_new,
+                Permission.Entity.PROJECT,
+                get_project_uuid(),
+            )
 
             if self.current_user.is_admin:
                 if self.owner:
@@ -187,6 +179,9 @@ class UserResource(DatabaseResource):
             'password',
             'password_confirmation',
             'password_current',
+            'owner',
+            'project_access',
+            'roles_display',
         ]), **kwargs)
 
     @safe_db_query

mage_ai/api/resources/WorkspaceResource.py

@@ -1,7 +1,8 @@
 import os
-import shutil
+import uuid
 from typing import Dict, List
 
+import ruamel.yaml
 import yaml
 
 from mage_ai.api.errors import ApiError
@@ -14,21 +15,19 @@ from mage_ai.cluster_manager.constants import (
     GCP_PROJECT_ID,
     GCP_REGION,
     KUBE_NAMESPACE,
-    KUBE_SERVICE_ACCOUNT_NAME,
-    KUBE_STORAGE_CLASS_NAME,
 )
 from mage_ai.data_preparation.repo_manager import (
     ProjectType,
     get_project_type,
     get_repo_config,
     get_repo_path,
-    init_repo,
 )
 from mage_ai.data_preparation.shared.constants import MANAGE_ENV_VAR
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Permission, Role, User
 from mage_ai.server.api.clusters import ClusterType
 from mage_ai.server.logger import Logger
+from mage_ai.settings import REQUIRE_USER_AUTHENTICATION
 
 logger = Logger().new_server_logger(__name__)
 
@@ -47,7 +46,8 @@ class WorkspaceResource(GenericResource):
         query_user = None
         if user_id:
             user_id = user_id[0]
-            query_user = User.query.get(user_id)
+            if user_id:
+                query_user = User.query.get(user_id)
 
         instances = self.get_instances(cluster_type)
         instance_map = {
@@ -60,7 +60,7 @@ class WorkspaceResource(GenericResource):
         repo_path = get_repo_path()
         projects_folder = os.path.join(repo_path, 'projects')
         if is_main_project:
-            projects = [f.name for f in os.scandir(projects_folder) if f.is_dir()]
+            projects = [f.name.split('.')[0] for f in os.scandir(projects_folder) if not f.is_dir()]
         else:
             projects = [instance.get('name') for instance in instances]
 
@@ -68,19 +68,25 @@ class WorkspaceResource(GenericResource):
         for project in projects:
             if project in instance_map:
                 try:
-                    repo_path = os.path.join(projects_folder, project)
                     workspace = dict(
                         name=project,
-                        repo_path=repo_path,
                         cluster_type=cluster_type,
                         instance=instance_map[project],
                     )
-                    if is_main_project and query_user:
-                        workspace['access'] = query_user.get_access(
-                            Permission.Entity.PROJECT,
-                            repo_path,
-                        )
+                    if is_main_project:
+                        workspace_file = os.path.join(projects_folder, f'{project}.yaml')
+                        config = {}
+                        with open(workspace_file) as f:
+                            config = yaml.full_load(f.read())
+                        project_uuid = config['project_uuid']
+                        workspace['project_uuid'] = project_uuid
+                        if query_user:
+                            workspace['access'] = query_user.get_access(
+                                Permission.Entity.PROJECT,
+                                project_uuid,
+                            )
                     workspaces.append(workspace)
+
                 except Exception as e:
                     print(f'Error fetching workspace: {str(e)}')
 
@@ -113,18 +119,33 @@ class WorkspaceResource(GenericResource):
         if not cluster_type:
             cluster_type = payload.get('cluster_type')
 
-        workspace_name = payload.get('name')
-
         error = ApiError.RESOURCE_ERROR.copy()
+        workspace_name = payload.pop('name')
+        if not workspace_name:
+            error.update(message='Please enter a valid workspace name.')
+            raise ApiError(error)
 
-        workspace_folder = None
-        if get_project_type() == ProjectType.MAIN:
-            workspace_folder = os.path.join(get_repo_path(), 'projects', workspace_name)
-            if os.path.exists(workspace_folder):
+        workspace_file = None
+        project_uuid = None
+        project_type = get_project_type()
+        if project_type == ProjectType.MAIN:
+            project_folder = os.path.join(get_repo_path(), 'projects')
+            if not os.path.exists(project_folder):
+                os.makedirs(project_folder)
+            workspace_file = os.path.join(project_folder, f'{workspace_name}.yaml')
+            if os.path.exists(workspace_file):
                 error.update(message=f'Project with name {workspace_name} already exists')
                 raise ApiError(error)
             try:
-                init_repo(workspace_folder, project_type=ProjectType.SUB)
+                yml = ruamel.yaml.YAML()
+                yml.preserve_quotes = True
+                yml.indent(mapping=2, sequence=2, offset=0)
+
+                project_uuid = uuid.uuid4().hex
+                data = dict(project_uuid=project_uuid)
+
+                with open(workspace_file, 'w') as f:
+                    yml.dump(data, f)
             except Exception as e:
                 error.update(message=f'Error creating project: {str(e)}')
                 raise ApiError(error)
@@ -133,27 +154,19 @@ class WorkspaceResource(GenericResource):
                 from mage_ai.cluster_manager.kubernetes.workload_manager import (
                     WorkloadManager,
                 )
-                namespace = payload.get('namespace', os.getenv(KUBE_NAMESPACE))
-                storage_class_name = payload.get(
-                    'storage_class_name',
-                    os.getenv(KUBE_STORAGE_CLASS_NAME),
-                )
-                service_account_name = payload.get(
-                    'service_account_name',
-                    os.getenv(KUBE_SERVICE_ACCOUNT_NAME),
-                )
-                container_config_yaml = payload.get('container_config')
-                container_config = None
-                if container_config_yaml:
-                    container_config = yaml.full_load(container_config_yaml)
+                namespace = payload.pop('namespace', os.getenv(KUBE_NAMESPACE))
 
                 k8s_workload_manager = WorkloadManager(namespace)
-                k8s_workload_manager.create_deployment(
+                extra_args = {}
+                if project_type == ProjectType.MAIN:
+                    extra_args = {
+                        'project_type': ProjectType.SUB,
+                        'project_uuid': project_uuid,
+                    }
+                k8s_workload_manager.create_workload(
                     workspace_name,
-                    container_config=container_config,
-                    service_account_name=service_account_name,
-                    storage_class_name=storage_class_name,
-                    volume_host_path=workspace_folder,
+                    **payload,
+                    **extra_args,
                 )
             elif cluster_type == ClusterType.ECS:
                 from mage_ai.cluster_manager.aws.ecs_task_manager import EcsTaskManager
@@ -193,15 +206,17 @@ class WorkspaceResource(GenericResource):
 
                 cloud_run_service_manager.create_service(workspace_name)
         except Exception as e:
-            if workspace_folder and os.path.exists(workspace_folder):
-                shutil.rmtree(workspace_folder)
+            if workspace_file and os.path.exists(workspace_file):
+                os.remove(workspace_file)
             error.update(message=str(e))
             raise ApiError(error)
 
-        if get_project_type() == ProjectType.MAIN:
+        if project_type == ProjectType.MAIN and \
+                project_uuid is not None and \
+                REQUIRE_USER_AUTHENTICATION:
             Role.create_default_roles(
                 entity=Permission.Entity.PROJECT,
-                entity_id=workspace_folder,
+                entity_id=project_uuid,
                 prefix=workspace_name,
             )
 
@@ -240,17 +255,32 @@ class WorkspaceResource(GenericResource):
         instance = self.model.get('instance')
 
         repo_path = get_repo_path()
-        project_folder = os.path.join(repo_path, 'projects', workspace_name)
+        workspace_file = os.path.join(repo_path, 'projects', f'{workspace_name}.yaml')
 
-        if get_project_type() == ProjectType.MAIN:
-            shutil.rmtree(project_folder)
-        if cluster_type == 'ecs':
-            from mage_ai.cluster_manager.aws.ecs_task_manager import EcsTaskManager
-            task_arn = instance.get('task_arn')
-            cluster_name = os.getenv(ECS_CLUSTER_NAME)
+        error = ApiError.RESOURCE_ERROR.copy()
 
-            ecs_instance_manager = EcsTaskManager(cluster_name)
-            ecs_instance_manager.delete_task(workspace_name, task_arn)
+        try:
+            if cluster_type == ClusterType.ECS:
+                from mage_ai.cluster_manager.aws.ecs_task_manager import EcsTaskManager
+                task_arn = instance.get('task_arn')
+                cluster_name = os.getenv(ECS_CLUSTER_NAME)
+
+                ecs_instance_manager = EcsTaskManager(cluster_name)
+                ecs_instance_manager.delete_task(workspace_name, task_arn)
+            elif cluster_type == ClusterType.K8S:
+                from mage_ai.cluster_manager.kubernetes.workload_manager import (
+                    WorkloadManager,
+                )
+                namespace = os.getenv(KUBE_NAMESPACE)
+
+                k8s_workload_manager = WorkloadManager(namespace)
+                k8s_workload_manager.delete_workload(workspace_name)
+        except Exception as e:
+            error.update(message=str(e))
+            raise ApiError(error)
+
+        if get_project_type() == ProjectType.MAIN:
+            os.remove(workspace_file)
 
         return self
 
@@ -265,7 +295,7 @@ class WorkspaceResource(GenericResource):
         if project_type == ProjectType.MAIN and subproject:
             repo_path = get_repo_path()
             projects_folder = os.path.join(repo_path, 'projects')
-            projects = [name for name in os.listdir(projects_folder) if os.path.isdir(name)]
+            projects = [f.name.split('.')[0] for f in os.scandir(projects_folder) if not f.is_dir()]
             if subproject not in projects:
                 error = ApiError.RESOURCE_NOT_FOUND.copy()
                 error.update(message=f'Project {subproject} was not found.')
@@ -284,7 +314,7 @@ class WorkspaceResource(GenericResource):
             namespace = os.getenv(KUBE_NAMESPACE)
             workload_manager = WorkloadManager(namespace)
 
-            instances = workload_manager.list_services()
+            instances = workload_manager.list_workloads()
         elif cluster_type == ClusterType.ECS:
             from mage_ai.cluster_manager.aws.ecs_task_manager import EcsTaskManager
 

mage_ai/authentication/oauth/active_directory.py

@@ -0,0 +1,17 @@
+from typing import Dict
+
+import requests
+
+
+def get_user_info(token: str) -> Dict:
+    url = 'https://graph.microsoft.com/v1.0/me'
+
+    headers = {
+        'Content-Type': 'application\\json',
+        'Authorization': 'Bearer {}'.format(token)
+    }
+
+    resp = requests.get(url, headers=headers)
+    result = resp.json()
+
+    return result

mage_ai/authentication/oauth/constants.py

@@ -1,3 +1,12 @@
+ACTIVE_DIRECTORY_CLIENT_ID = '51aec820-9d49-40a9-b046-17c1f28f620d'
+
 GITHUB_CLIENT_ID = '8577f13ddc81e2848b07'
 GITHUB_STATE = '1337'
+
+OAUTH_PROVIDER_ACTIVE_DIRECTORY = 'active_directory'
 OAUTH_PROVIDER_GITHUB = 'github'
+
+VALID_OAUTH_PROVIDERS = [
+    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
+    OAUTH_PROVIDER_GITHUB,
+]

mage_ai/authentication/oauth/utils.py

@@ -1,7 +1,8 @@
 from datetime import datetime
-from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
 from typing import List
 
+from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
+
 
 def access_tokens_for_provider(provider: str) -> List[Oauth2AccessToken]:
     oauth_client = Oauth2Application.query.filter(

mage_ai/authentication/oauth2.py

@@ -1,9 +1,15 @@
+import os
+import secrets
 from datetime import datetime, timedelta
-from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application, User
 from typing import Dict
+
 import jwt
-import os
-import secrets
+
+from mage_ai.orchestration.db.models.oauth import (
+    Oauth2AccessToken,
+    Oauth2Application,
+    User,
+)
 
 JWT_ALGORITHM = 'HS256'
 JWT_SECRET = os.getenv('JWT_SECRET', 'materia')