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

mage_ai/api/policies/GlobalHookPolicy.py

@@ -0,0 +1,126 @@
+from mage_ai.api.oauth_scope import OauthScope
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.api.policies.BasePolicy import BasePolicy
+from mage_ai.api.presenters.GlobalHookPresenter import GlobalHookPresenter
+
+WRITEABLE_ATTRIBUTES = [
+    'conditions',
+    'operation_type',
+    'outputs',
+    'pipeline',
+    'predicates',
+    'resource_type',
+    'run_settings',
+    'stages',
+    'strategies',
+]
+
+
+class GlobalHookPolicy(BasePolicy):
+    pass
+
+
+GlobalHookPolicy.allow_actions(
+    [
+        OperationType.CREATE,
+        OperationType.DELETE,
+        OperationType.DETAIL,
+        OperationType.LIST,
+        OperationType.UPDATE,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_admin_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+GlobalHookPolicy.allow_read(
+    GlobalHookPresenter.default_attributes + [
+        'operation_types',
+        'pipeline_details',
+        'resource_types',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        OperationType.CREATE,
+        OperationType.DELETE,
+        OperationType.DETAIL,
+        OperationType.LIST,
+        OperationType.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_admin_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+GlobalHookPolicy.allow_write(
+    WRITEABLE_ATTRIBUTES + [
+        'uuid',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        OperationType.CREATE,
+    ],
+    condition=lambda policy: policy.has_at_least_admin_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+GlobalHookPolicy.allow_write(
+    WRITEABLE_ATTRIBUTES + [
+        'snapshot',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        OperationType.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_admin_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+GlobalHookPolicy.allow_query(
+    [
+        'include_snapshot_validation',
+        'operation_type[]',
+        'resource_type[]',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        OperationType.LIST,
+    ],
+    condition=lambda policy: policy.has_at_least_admin_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+GlobalHookPolicy.allow_query(
+    [
+        'include_operation_types',
+        'include_resource_types',
+        'include_snapshot_validation',
+        'operation_type',
+        'resource_type',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        OperationType.CREATE,
+        OperationType.DELETE,
+        OperationType.DETAIL,
+        OperationType.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_admin_role(),
+    override_permission_condition=lambda _policy: True,
+)

mage_ai/api/policies/OauthPolicy.py

@@ -1,11 +1,60 @@
+from typing import Dict
+
+from mage_ai.api.constants import AttributeOperationType, AttributeType
 from mage_ai.api.oauth_scope import OauthScope
 from mage_ai.api.operations import constants
+from mage_ai.api.operations.constants import OperationType
 from mage_ai.api.policies.BasePolicy import BasePolicy
+from mage_ai.api.policies.mixins.user_permissions import UserPermissionMixIn
 from mage_ai.api.presenters.OauthPresenter import OauthPresenter
-
-
-class OauthPolicy(BasePolicy):
-    pass
+from mage_ai.shared.hash import merge_dict
+
+
+class OauthPolicy(BasePolicy, UserPermissionMixIn):
+    @classmethod
+    def action_rule_with_permissions(self, operation: OperationType) -> Dict:
+        return merge_dict(
+            super().action_rule_with_permissions(operation),
+            {
+                OauthScope.CLIENT_PUBLIC: [
+                    dict(
+                        condition=lambda _policy: operation
+                        in [OperationType.DETAIL, OperationType.LIST],
+                    ),
+                ],
+            },
+        )
+
+    @classmethod
+    def attribute_rule_with_permissions(
+        self,
+        attribute_operation_type: AttributeOperationType,
+        resource_attribute: str,
+    ) -> Dict:
+        config = {}
+        if AttributeOperationType.READ == attribute_operation_type:
+            config = self.read_rules[self.__name__].get(resource_attribute)
+        elif AttributeOperationType.QUERY == attribute_operation_type:
+            config = self.query_rules[self.__name__].get(AttributeType.ALL)
+        else:
+            config = self.write_rules[self.__name__].get(resource_attribute)
+
+        return merge_dict(
+            super().attribute_rule_with_permissions(
+                attribute_operation_type,
+                resource_attribute,
+            ),
+            {
+                OauthScope.CLIENT_PUBLIC: {
+                    OperationType.DETAIL: config[OauthScope.CLIENT_PUBLIC][
+                        OperationType.DETAIL
+                    ],
+                    OperationType.LIST: config[OauthScope.CLIENT_PUBLIC][
+                        OperationType.LIST
+                    ],
+                },
+            },
+        )
 
 
 OauthPolicy.allow_actions(
@@ -23,6 +72,7 @@ OauthPolicy.allow_actions(
 OauthPolicy.allow_actions(
     [
         constants.DETAIL,
+        constants.LIST,
     ],
     scopes=[
         OauthScope.CLIENT_PRIVATE,
@@ -52,6 +102,7 @@ OauthPolicy.allow_read(
         OauthScope.CLIENT_PUBLIC,
     ],
     on_action=[
+        constants.LIST,
         constants.DETAIL,
     ],
     condition=lambda policy: policy.has_at_least_viewer_role(),
@@ -89,6 +140,7 @@ OauthPolicy.allow_write(
 
 OauthPolicy.allow_query(
     on_action=[
+        constants.LIST,
         constants.DETAIL,
     ],
     scopes=[

mage_ai/api/policies/ProjectPolicy.py

@@ -40,6 +40,7 @@ ProjectPolicy.allow_write([
     'emr_config',
     'features',
     'openai_api_key',
+    'pipelines',
     'remote_variables_dir',
     'spark_config',
 ], scopes=[

mage_ai/api/presenters/ComputeClusterPresenter.py

@@ -0,0 +1,19 @@
+from typing import Dict
+
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+from mage_ai.shared.models import BaseDataClass
+
+
+class ComputeClusterPresenter(BasePresenter):
+    default_attributes = [
+        'cluster',
+    ]
+
+    async def prepare_present(self, **kwargs) -> Dict:
+        data = self.resource.model
+        if 'cluster' in data:
+            cluster = data.get('cluster')
+            if isinstance(cluster, BaseDataClass):
+                data['cluster'] = cluster.to_dict()
+
+        return data

mage_ai/api/presenters/ComputeConnectionPresenter.py

@@ -0,0 +1,26 @@
+from typing import Dict
+
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class ComputeConnectionPresenter(BasePresenter):
+    default_attributes = [
+        'actions',
+        'attributes',
+        'connection',
+        'description',
+        'error',
+        'group',
+        'name',
+        'required',
+        'state',
+        'status',
+        'status_calculated',
+        'steps',
+        'tab',
+        'target',
+        'uuid',
+    ]
+
+    async def prepare_present(self, **kwargs) -> Dict:
+        return self.resource.model.to_dict()

mage_ai/api/presenters/ComputeServicePresenter.py

@@ -0,0 +1,15 @@
+from typing import Dict
+
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class ComputeServicePresenter(BasePresenter):
+    default_attributes = [
+        'connection_credentials',
+        'clusters',
+        'setup_steps',
+        'uuid',
+    ]
+
+    async def prepare_present(self, **kwargs) -> Dict:
+        return self.resource.model.to_dict()

mage_ai/api/presenters/DownloadPresenter.py

@@ -0,0 +1,13 @@
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class DownloadPresenter(BasePresenter):
+    default_attributes = [
+        'uri',
+    ]
+
+    def present(self, **kwargs):
+        if type(self.model) is dict:
+            return self.model
+
+        return self.model.to_dict()

mage_ai/api/presenters/GlobalHookPresenter.py

@@ -0,0 +1,74 @@
+from typing import Dict
+
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+from mage_ai.data_preparation.models.global_hooks.constants import RESOURCE_TYPES
+from mage_ai.data_preparation.models.global_hooks.models import HookOperation
+
+
+class GlobalHookPresenter(BasePresenter):
+    default_attributes = [
+        'conditions',
+        'metadata',
+        'operation_type',
+        'outputs',
+        'pipeline',
+        'predicates',
+        'resource_type',
+        'run_settings',
+        'stages',
+        'strategies',
+        'uuid',
+    ]
+
+    async def prepare_present(self, **kwargs) -> Dict:
+        query = kwargs.get('query') or {}
+
+        data = {}
+        if self.resource.model:
+            include_snapshot_validation = False
+            if 'include_snapshot_validation' in query:
+                include_snapshot_validation = True
+
+            data = self.resource.model.to_dict(
+                include_all=True,
+                include_snapshot_validation=include_snapshot_validation,
+            )
+
+        display_format = kwargs.get('format')
+        if 'with_pipeline_details' == display_format:
+            data['pipeline_details'] = await self.pipeline_details(**kwargs)
+
+        if 'include_operation_types' in query:
+            data['operation_types'] = [m.value for m in HookOperation]
+
+        if 'include_resource_types' in query:
+            data['resource_types'] = RESOURCE_TYPES
+
+        return data
+
+    async def pipeline_details(self, **kwargs) -> Dict:
+        if self.resource.model.pipeline_settings and \
+                self.resource.model.pipeline_settings.get('uuid'):
+
+            pipeline = await self.resource.pipeline
+            pipeline_dict = await pipeline.to_dict_async()
+            return pipeline_dict
+
+
+GlobalHookPresenter.register_format(
+    'with_pipeline_details',
+    GlobalHookPresenter.default_attributes + [
+        'pipeline_details',
+    ],
+)
+
+
+GlobalHookPresenter.register_format(
+    OperationType.DETAIL,
+    GlobalHookPresenter.default_attributes + [
+        'operation_types',
+        'pipeline_details',
+        'resource_types',
+    ],
+)

mage_ai/api/presenters/PipelinePresenter.py

@@ -22,6 +22,7 @@ class PipelinePresenter(BasePresenter):
         'notification_config',
         'retry_config',
         'run_pipeline_in_one_process',
+        'settings',
         'spark_config',
         'tags',
         'type',

mage_ai/api/presenters/ProjectPresenter.py

@@ -9,6 +9,7 @@ class ProjectPresenter(BasePresenter):
         'latest_version',
         'name',
         'openai_api_key',
+        'pipelines',
         'project_uuid',
         'remote_variables_dir',
         'spark_config',

mage_ai/api/presenters/SparkApplicationPresenter.py

@@ -4,6 +4,8 @@ from mage_ai.api.presenters.BasePresenter import BasePresenter
 class SparkApplicationPresenter(BasePresenter):
     default_attributes = [
         'attempts',
+        'attempts_count',
+        'calculated_id',
         'id',
         'name',
         'spark_ui_url',

mage_ai/api/presenters/WorkspacePresenter.py

@@ -1,30 +1,30 @@
+from dataclasses import fields
+
 from mage_ai.api.presenters.BasePresenter import BasePresenter
+from mage_ai.cluster_manager.config import (
+    CloudRunWorkspaceConfig,
+    EcsWorkspaceConfig,
+    KubernetesWorkspaceConfig,
+)
 from mage_ai.shared.hash import merge_dict
 
 
 class WorkspacePresenter(BasePresenter):
-    default_attributes = [
-        'access',
-        'cluster_name',
-        'cluster_type',
-        'container_config',
-        'container_name',
-        'instance',
-        'lifecycle_config',
-        'name',
-        'namespace',
-        'path_to_credentials',
-        'project_uuid',
-        'project_id',
-        'region',
-        'repo_path',
-        'service_account_name',
-        'storage_class_name',
-        'storage_access_mode',
-        'storage_request_size',
-        'success',
-        'task_definition',
-    ]
+    default_attributes = list(
+        set(
+            [
+                'access',
+                'cluster_type',
+                'instance',
+                'name',
+                'repo_path',
+                'success',
+                *[f.name for f in fields(KubernetesWorkspaceConfig)],
+                *[f.name for f in fields(CloudRunWorkspaceConfig)],
+                *[f.name for f in fields(EcsWorkspaceConfig)],
+            ]
+        )
+    )
 
     async def present(self, **kwargs):
         workspace = self.model.pop('workspace', None)

mage_ai/api/resources/AsyncBaseResource.py

@@ -0,0 +1,39 @@
+from mage_ai.api.resources.BaseResource import BaseResource
+from mage_ai.api.result_set import ResultSet
+
+
+class AsyncBaseResource(BaseResource):
+    async def collective_load_for_attribute(self, key):
+        k_name = self.__class__.__name__
+        if self.result_set().context and self.result_set().context.data:
+            loaded = self.result_set().context.data.get(k_name, {}).get(key, None)
+        else:
+            loaded = None
+        loader = self.__class__.collective_loader().get(key, None)
+        if not loaded and loader:
+            loaded = await loader['load'](self)
+            if loaded and not isinstance(
+                    loaded,
+                    ResultSet) and not isinstance(
+                    loaded,
+                    dict):
+                loaded = ResultSet(loaded)
+            if not self.result_set().context.data.get(k_name):
+                self.result_set().context.data[k_name] = {}
+            self.result_set().context.data[k_name][key] = loaded
+        return loaded
+
+    async def __getattr__(self, name):
+        async def _missing(*args, **kwargs):
+            loader = self.__class__.collective_loader().get(name, None)
+            if loader:
+                arr = await self.collective_load_for_attribute(name)
+                if loader['select']:
+                    val = await loader['select'](self, arr)
+                else:
+                    val = arr
+            else:
+                val = getattr(self.model, name)
+
+            return val
+        return await _missing()

mage_ai/api/resources/ClusterResource.py

@@ -41,7 +41,7 @@ class ClusterResource(GenericResource):
                 error['message'] = 'Please include cluster_id in the request payload.'
                 raise ApiError(error)
 
-            emr_cluster_manager.set_active_cluster(cluster_id)
+            emr_cluster_manager.set_active_cluster(cluster_id=cluster_id)
             success = True
 
             self.model.update(

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