mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

mage_ai/api/resources/ComputeClusterResource.py

@@ -0,0 +1,109 @@
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.services.compute.aws.models import Cluster
+from mage_ai.services.compute.models import ComputeService
+from mage_ai.services.spark.constants import ComputeServiceUUID
+from mage_ai.services.ssh.aws.emr.models import close_tunnel
+from mage_ai.services.ssh.aws.emr.utils import tunnel
+
+
+class ComputeClusterResource(GenericResource):
+    @classmethod
+    async def collection(self, query_arg, _meta, user, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        include_all_states = query_arg.get('include_all_states', [None])
+        if include_all_states:
+            include_all_states = include_all_states[0]
+
+        clusters = []
+
+        if isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                try:
+                    result = parent_model.clusters_and_metadata(
+                        include_all_states=include_all_states,
+                    )
+                    clusters = [dict(
+                        cluster=cluster,
+                    ) for cluster in result.get('clusters') or []]
+                except Exception as err:
+                    print(f'[WARNING] ComputeClusterResource.collection: {err}')
+
+        return self.build_result_set(
+            clusters,
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def member(self, pk, user, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        cluster = None
+
+        if isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                cluster = parent_model.get_cluster_details(pk)
+
+        return self(dict(
+            cluster=cluster,
+        ), user, **kwargs)
+
+    @classmethod
+    def create(self, payload, user, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        cluster = None
+
+        if isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                cluster = parent_model.create_cluster()
+
+        return self(dict(
+            cluster=cluster,
+        ), user, **kwargs)
+
+    async def delete(self, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        if isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                cluster = self.get_cluster()
+                parent_model.terminate_clusters([cluster.id])
+
+                if cluster.active:
+                    close_tunnel()
+
+    async def update(self, payload, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        if isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                cluster_id = self.get_cluster_id()
+                if cluster_id:
+                    cluster = parent_model.update_cluster(cluster_id, payload)
+                    if cluster:
+                        self.model = dict(cluster=cluster)
+
+                        if cluster.active:
+                            def _callback(*args, **kwargs):
+                                tunnel(reconnect=True)
+
+                            self.on_update_callback = _callback
+
+    def get_cluster(self) -> Cluster:
+        return self.model.get('cluster')
+
+    def get_cluster_id(self) -> str:
+        cluster_id = None
+
+        if 'cluster' in self.model:
+            cluster = self.get_cluster()
+            if cluster:
+                cluster_id = None
+                if isinstance(cluster, Cluster):
+                    cluster_id = cluster.id
+                elif isinstance(cluster, dict):
+                    cluster_id = cluster.get('id')
+
+        return cluster_id

mage_ai/api/resources/ComputeConnectionResource.py

@@ -0,0 +1,103 @@
+from mage_ai.api.errors import ApiError
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.services.compute.constants import (
+    ComputeConnectionActionUUID,
+    ComputeConnectionState,
+)
+from mage_ai.services.compute.models import ComputeConnection, ComputeService
+from mage_ai.services.spark.constants import ComputeServiceUUID
+from mage_ai.services.ssh.aws.emr.models import SSHTunnel, create_tunnel
+from mage_ai.services.ssh.aws.emr.utils import should_tunnel, tunnel
+from mage_ai.shared.array import find
+
+
+class ComputeConnectionResource(GenericResource):
+    @classmethod
+    async def collection(self, query_arg, _meta, user, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        arr = []
+
+        if parent_model and isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                arr.extend(parent_model.compute_connections())
+
+        return self.build_result_set(
+            arr,
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def member(self, pk, user, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        model = ComputeConnection.load(name=pk, uuid=pk)
+        if parent_model and isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                model = find(lambda x: x.uuid == pk, parent_model.compute_connections())
+
+        return self(model, user, **kwargs)
+
+    async def update(self, payload, **kwargs):
+        parent_model = kwargs.get('parent_model')
+
+        action_uuid = payload.get('action')
+
+        if not action_uuid:
+            return
+
+        if parent_model and isinstance(parent_model, ComputeService):
+            if ComputeServiceUUID.AWS_EMR == parent_model.uuid:
+                from mage_ai.services.compute.aws.steps import SetupStepUUID
+
+                if SetupStepUUID.OBSERVABILITY == self.model.uuid:
+                    action_valid = action_uuid in [
+                        ComputeConnectionActionUUID.CREATE,
+                        ComputeConnectionActionUUID.DELETE,
+                        ComputeConnectionActionUUID.DESELECT,
+                        ComputeConnectionActionUUID.UPDATE,
+                    ]
+
+                    if action_valid:
+                        if ComputeConnectionActionUUID.CREATE == action_uuid:
+                            try:
+                                should_tunnel(
+                                    ignore_active_kernel=True,
+                                    raise_error=True,
+                                )
+
+                                ssh_tunnel = create_tunnel(
+                                    connect=False,
+                                    fresh=True,
+                                    reconnect=False,
+                                    stop=False,
+                                )
+                                if ssh_tunnel:
+                                    ssh_tunnel.precheck_access(raise_error=True)
+                                    ssh_tunnel.connect()
+                                else:
+                                    raise Exception('Cannot create SSH tunnel instance.')
+                            except Exception as err:
+                                error = ApiError.UNAUTHORIZED_ACCESS.copy()
+                                error.update(message=str(err))
+                                raise ApiError(error)
+
+                        def _callback(
+                            action_uuid=action_uuid,
+                            *args,
+                            **kwargs,
+                        ):
+                            ssh_tunnel = SSHTunnel()
+
+                            if ComputeConnectionActionUUID.CREATE == action_uuid:
+                                tunnel(ignore_active_kernel=True)
+                            elif ComputeConnectionActionUUID.DELETE == action_uuid:
+                                ssh_tunnel.close()
+                            elif ComputeConnectionActionUUID.DESELECT == action_uuid:
+                                ssh_tunnel.stop()
+                            elif ComputeConnectionActionUUID.UPDATE == action_uuid:
+                                ssh_tunnel.reconnect()
+
+                        self.model.state = ComputeConnectionState.PENDING
+                        self.on_update_callback = _callback

mage_ai/api/resources/ComputeServiceResource.py

@@ -0,0 +1,35 @@
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.services.compute.models import ComputeService
+
+
+class ComputeServiceResource(GenericResource):
+    @classmethod
+    async def collection(self, _query, _meta, user, **kwargs):
+        resource = await self.member(None, user, **kwargs)
+
+        return self.build_result_set(
+            [
+                resource.model,
+            ],
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    def get_model(self, _pk):
+        return ComputeService.build(Project())
+
+    @classmethod
+    async def member(self, pk, user, **kwargs):
+        query_arg = kwargs.get('query', {})
+
+        with_clusters = query_arg.get('with_clusters', [False])
+        if with_clusters:
+            with_clusters = with_clusters[0]
+
+        compute_service = self.get_model(pk)
+        if with_clusters:
+            compute_service.with_clusters = with_clusters
+
+        return self(compute_service, user, **kwargs)

mage_ai/api/resources/DownloadResource.py

@@ -0,0 +1,56 @@
+import datetime
+from typing import Dict
+
+import jwt
+
+from mage_ai.api.constants import DOWNLOAD_TOKEN_LIFESPAN
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.data_preparation.models.download import Download
+from mage_ai.data_preparation.models.file import File
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.orchestration.db import safe_db_query
+from mage_ai.settings import JWT_DOWNLOAD_SECRET
+
+
+class DownloadResource(GenericResource):
+
+    @classmethod
+    @safe_db_query
+    async def create(self, payload: Dict, user, **kwargs) -> 'DownloadResource':
+        parent_model = kwargs.get('parent_model')
+
+        ignore_folder_structure = payload.get('ignore_folder_structure', False)
+
+        file_list = []
+
+        if isinstance(parent_model, Pipeline):
+            file_name, file_list = self.fetch_pipeline_files(parent_model)
+        elif isinstance(parent_model, File):
+            if not parent_model.exists():
+                raise Exception(f'File {parent_model.filename} does not exist.')
+            file_name = parent_model.filename
+            file_list.append(parent_model.file_path)
+
+        token = self.generate_download_token(file_name, file_list, ignore_folder_structure)
+        download = Download(token)
+
+        return self(download, user)
+
+    def fetch_pipeline_files(pipeline: Pipeline):
+        config_path = pipeline.config_path
+
+        zip_name = f'{pipeline.uuid}.zip'
+        file_list = [config_path] + [block.file_path for block in pipeline.blocks_by_uuid.values()]
+
+        return [zip_name, file_list]
+
+    def generate_download_token(file_name, file_list, ignore_folder_structure):
+        now = datetime.datetime.utcnow()
+        payload = {
+            'file_name': file_name,
+            'file_list': file_list,
+            'ignore_folder_structure': ignore_folder_structure,
+            'iat': now,
+            'exp': now + datetime.timedelta(seconds=DOWNLOAD_TOKEN_LIFESPAN)
+        }
+        return jwt.encode(payload, JWT_DOWNLOAD_SECRET, algorithm="HS256")

mage_ai/api/resources/ExecutionStateResource.py

@@ -18,7 +18,7 @@ class ExecutionStateResource(GenericResource):
 
         if pipeline_uuid and block_uuid:
             pipeline = await Pipeline.get_async(pipeline_uuid)
-            if pipeline and PipelineType.PYTHON == pipeline.type:
+            if pipeline and pipeline.type in [PipelineType.PYTHON, PipelineType.PYSPARK]:
                 block = pipeline.get_block(block_uuid)
                 if block and block.compute_management_enabled() and block.is_using_spark():
                     arr.append(dict(

mage_ai/api/resources/GitBranchResource.py

@@ -109,6 +109,15 @@ class GitBranchResource(GenericResource):
         message = payload.get('message', None)
         remote = payload.get('remote', None)
 
+        access_token = api.get_access_token_for_user(self.current_user)
+        http_access_token = git_manager.get_access_token()
+
+        token = None
+        if access_token:
+            token = access_token.token
+        elif http_access_token:
+            token = http_access_token
+
         if action_type == 'status':
             status = git_manager.status()
             untracked_files = await git_manager.untracked_files()
@@ -134,14 +143,15 @@ class GitBranchResource(GenericResource):
                 from git.exc import GitCommandError
 
                 try:
-                    access_token = api.get_access_token_for_user(self.current_user)
-                    if access_token:
-                        remote = git_manager.repo.remotes[push['remote']]
+                    remote = git_manager.repo.remotes[push['remote']]
+                    url = list(remote.urls)[0]
+
+                    if token:
                         custom_progress = api.push(
                             remote.name,
-                            [url for url in remote.urls][0],
+                            url,
                             push['branch'],
-                            access_token.token,
+                            token,
                         )
                         if custom_progress.other_lines:
                             lines = custom_progress.other_lines
@@ -161,25 +171,24 @@ class GitBranchResource(GenericResource):
                 from git.exc import GitCommandError
 
                 try:
-                    access_token = api.get_access_token_for_user(self.current_user)
-                    if access_token:
-                        branch_name = pull.get('branch')
-                        remote = git_manager.repo.remotes[pull['remote']]
-                        url = list(remote.urls)[0]
+                    branch_name = pull.get('branch')
+                    remote = git_manager.repo.remotes[pull['remote']]
+                    url = list(remote.urls)[0]
 
+                    if token:
                         custom_progress = None
                         if branch_name:
                             custom_progress = api.pull(
                                 remote.name,
                                 url,
                                 pull.get('branch'),
-                                access_token.token,
+                                token,
                             )
                         else:
                             custom_progress = api.fetch(
                                 remote.name,
                                 url,
-                                access_token.token,
+                                token,
                             )
 
                         if custom_progress and custom_progress.other_lines:
@@ -200,15 +209,14 @@ class GitBranchResource(GenericResource):
                 from git.exc import GitCommandError
 
                 try:
-                    access_token = api.get_access_token_for_user(self.current_user)
-                    if access_token:
-                        remote = git_manager.repo.remotes[fetch['remote']]
-                        url = list(remote.urls)[0]
+                    remote = git_manager.repo.remotes[fetch['remote']]
+                    url = list(remote.urls)[0]
 
+                    if token:
                         custom_progress = api.fetch(
                             remote.name,
                             url,
-                            access_token.token,
+                            token,
                         )
 
                         if custom_progress and custom_progress.other_lines:
@@ -218,7 +226,7 @@ class GitBranchResource(GenericResource):
                             self.model['progress'] = lines
                     else:
                         self.model['error'] = \
-                            'Please authenticate with GitHub before trying to pull.'
+                            'Please authenticate with GitHub before trying to fetch.'
                 except GitCommandError as err:
                     self.model['error'] = str(err)
             else:
@@ -233,17 +241,16 @@ class GitBranchResource(GenericResource):
                     from git.exc import GitCommandError
 
                     try:
-                        access_token = api.get_access_token_for_user(self.current_user)
-                        if access_token:
-                            branch_name = reset.get('branch')
-                            remote = git_manager.repo.remotes[reset['remote']]
-                            url = list(remote.urls)[0]
+                        branch_name = reset.get('branch')
+                        remote = git_manager.repo.remotes[reset['remote']]
+                        url = list(remote.urls)[0]
 
+                        if token:
                             api.reset_hard(
                                 remote.name,
                                 url,
                                 branch_name,
-                                access_token.token,
+                                token,
                             )
                         else:
                             self.model['error'] = \
@@ -255,15 +262,14 @@ class GitBranchResource(GenericResource):
             if clone and 'remote' in clone:
                 from git.exc import GitCommandError
                 try:
-                    access_token = api.get_access_token_for_user(self.current_user)
-                    if access_token:
-                        remote = git_manager.repo.remotes[clone['remote']]
-                        url = list(remote.urls)[0]
+                    remote = git_manager.repo.remotes[clone['remote']]
+                    url = list(remote.urls)[0]
 
+                    if token:
                         api.clone(
                             remote.name,
                             url,
-                            access_token.token,
+                            token,
                         )
                     else:
                         self.model['error'] = \

mage_ai/api/resources/GitCustomBranchResource.py

@@ -10,6 +10,15 @@ class GitCustomBranchResource(GitBranchResource):
     def get_git_manager(self, user, setup_repo: bool = False) -> Git:
         return Git.get_manager(auth_type=AuthType.OAUTH, setup_repo=False, user=user)
 
+    @classmethod
+    async def member(self, pk, user, **kwargs):
+        resource = await GitBranchResource.member(pk, user, **kwargs)
+        model = resource.model
+        model['access_token_exists'] = (
+            self.get_git_manager(user=user).get_access_token() is not None
+        )
+        return self(model, user, **kwargs)
+
     def remotes(self, limit: int = None) -> List[Dict]:
         git_manager = self.get_git_manager(user=self.current_user)
         remotes = git_manager.remotes(limit=limit, user=self.current_user)

mage_ai/api/resources/GlobalHookResource.py

@@ -0,0 +1,192 @@
+import asyncio
+from typing import Dict, List
+
+from mage_ai.api.errors import ApiError
+from mage_ai.api.resources.AsyncBaseResource import AsyncBaseResource
+from mage_ai.api.resources.PipelineResource import PipelineResource
+from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.data_preparation.models.global_hooks.constants import (
+    DISABLED_RESOURCE_TYPES,
+)
+from mage_ai.data_preparation.models.global_hooks.models import (
+    GlobalHooks,
+    Hook,
+    HookOperation,
+)
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.shared.array import find
+from mage_ai.shared.hash import ignore_keys
+
+
+class GlobalHookResource(AsyncBaseResource):
+    @classmethod
+    async def collection(self, query: Dict, _meta: Dict, user: User, **kwargs):
+        operation_types = query.get('operation_type[]', [])
+        if operation_types:
+            operation_types = operation_types[0]
+        if isinstance(operation_types, str):
+            operation_types = [HookOperation(m) for m in operation_types.split(',') if m]
+
+        resource_types = query.get('resource_type[]', [])
+        if resource_types:
+            resource_types = resource_types[0]
+        if isinstance(resource_types, str):
+            resource_types = [EntityName(m) for m in resource_types.split(',') if m]
+
+        global_hooks = GlobalHooks.load_from_file()
+
+        return self.build_result_set(
+            global_hooks.hooks(
+                operation_types=operation_types,
+                resource_types=resource_types,
+            ),
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def create(self, payload: Dict, user: User, **kwargs) -> 'GlobalHookResource':
+        if user and user.id:
+            payload['metadata'] = dict(
+                user=dict(
+                    id=user.id,
+                ),
+            )
+
+        hook = Hook.load(**payload)
+
+        if not hook.uuid or not hook.resource_type or not hook.operation_type:
+            missing = []
+            if not hook.resource_type:
+                missing.append('resource_type')
+            if not hook.operation_type:
+                missing.append('operation_type')
+            if not hook.uuid:
+                missing.append('uuid')
+
+            error = ApiError.RESOURCE_INVALID.copy()
+            error.update(
+                message=f'Hook is missing the following required attributes: {", ".join(missing)}.',
+            )
+            raise ApiError(error)
+
+        if hook.resource_type and hook.resource_type in DISABLED_RESOURCE_TYPES:
+            error = ApiError.RESOURCE_INVALID.copy()
+            error.update(
+                message=f'Hooks cannot be created for resource type {hook.resource_type.value}.',
+            )
+            raise ApiError(error)
+
+        global_hooks = GlobalHooks.load_from_file()
+        if global_hooks.get_hook(
+            operation_type=hook.operation_type,
+            resource_type=hook.resource_type,
+            uuid=hook.uuid,
+        ):
+            error = ApiError.RESOURCE_INVALID.copy()
+            error.update(
+                message=f'Hook {hook.uuid} already exists for resource '
+                        f'{hook.resource_type} and operation {hook.operation_type}.',
+            )
+            raise ApiError(error)
+
+        global_hooks.add_hook(hook, snapshot=True)
+        global_hooks.save()
+
+        return self(hook, user, **kwargs)
+
+    @classmethod
+    async def member(self, pk: str, user: User, **kwargs):
+        query = kwargs.get('query') or {}
+
+        resource_type = query.get('resource_type', [None])
+        if resource_type:
+            resource_type = resource_type[0]
+
+        operation_type = query.get('operation_type', [None])
+        if operation_type:
+            operation_type = operation_type[0]
+
+        include_operation_types = query.get('include_operation_types', [None])
+        if include_operation_types:
+            include_operation_types = include_operation_types[0]
+
+        include_resource_types = query.get('include_resource_types', [None])
+        if include_resource_types:
+            include_resource_types = include_resource_types[0]
+
+        global_hooks = GlobalHooks.load_from_file()
+        hook = global_hooks.get_hook(
+            operation_type=HookOperation(operation_type) if operation_type else operation_type,
+            resource_type=EntityName(resource_type) if resource_type else resource_type,
+            uuid=pk,
+        )
+
+        if not hook and not include_operation_types and not include_resource_types:
+            raise ApiError(ApiError.RESOURCE_NOT_FOUND)
+
+        return self(hook, user, **kwargs)
+
+    async def update(self, payload: Dict, **kwargs):
+        global_hooks = GlobalHooks.load_from_file()
+        self.model = global_hooks.add_hook(
+            self.model,
+            payload=ignore_keys(payload, ['snapshot']),
+            snapshot=payload.get('snapshot') or False,
+            update=True,
+        )
+        global_hooks.save()
+
+    async def delete(self, **kwargs):
+        global_hooks = GlobalHooks.load_from_file()
+        global_hooks.remove_hook(self.model)
+        global_hooks.save()
+
+
+async def __load_pipelines(resource: GlobalHookResource):
+    pipeline_uuids = []
+    for res in resource.result_set():
+        if res.model.pipeline_settings and res.model.pipeline_settings.get('uuid'):
+            pipeline_uuid = res.model.pipeline_settings.get('uuid')
+            if pipeline_uuid:
+                pipeline_uuids.append(pipeline_uuid)
+
+    async def get_pipeline(uuid):
+        try:
+            return await Pipeline.get_async(uuid)
+        except Exception as err:
+            err_message = f'Error loading pipeline {uuid}: {err}.'
+            if err.__class__.__name__ == 'OSError' and 'Too many open files' in err.strerror:
+                raise Exception(err_message)
+            else:
+                print(err_message)
+                return None
+
+    pipelines = await asyncio.gather(
+        *[get_pipeline(uuid) for uuid in pipeline_uuids]
+    )
+    pipelines = [p for p in pipelines if p is not None]
+
+    return PipelineResource.build_result_set(pipelines, resource.current_user)
+
+
+async def __find_pipeline(resource: GlobalHookResource, arr: List[PipelineResource]):
+    pipeline_uuid = resource.model.pipeline_settings and resource.model.pipeline_settings.get(
+        'uuid',
+    )
+
+    if not pipeline_uuid:
+        return
+
+    return find(
+        lambda x, pipeline_uuid=pipeline_uuid: x.uuid == pipeline_uuid,
+        arr,
+    )
+
+
+GlobalHookResource.register_collective_loader(
+    'pipeline_details',
+    load=__load_pipelines,
+    select=__find_pipeline,
+)

mage_ai/api/resources/KernelResource.py

@@ -7,6 +7,7 @@ from mage_ai.server.active_kernel import (
     switch_active_kernel,
 )
 from mage_ai.server.kernels import DEFAULT_KERNEL_NAME, KernelName, kernel_managers
+from mage_ai.services.ssh.aws.emr.utils import tunnel
 
 
 class KernelResource(GenericResource):
@@ -56,4 +57,13 @@ class KernelResource(GenericResource):
                 if 'start_kernel' in str(e):
                     start_kernel()
 
+        def _callback(*args, **kwargs):
+            tunnel(
+                kernel_name=self.model.kernel_name,
+                reconnect=True,
+                validate_conditions=True,
+            )
+
+        self.on_update_callback = _callback
+
         return self