mage-ai 0.9.44__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(),
@@ -88,10 +139,8 @@ OauthPolicy.allow_write(
 
 
 OauthPolicy.allow_query(
-    [
-        'redirect_uri',
-    ],
     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/policies/RolePolicy.py

@@ -50,10 +50,21 @@ RolePolicy.allow_read(RolePresenter.default_attributes + [
 ], on_action=[
     constants.CREATE,
     constants.DELETE,
-    constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_admin_role())
 
 
+# allow all logged in users to view roles on update because they can update
+# their own profile, and roles will be passed back in the response.
+RolePolicy.allow_read(RolePresenter.default_attributes + [
+    'user_id',
+    'users',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.UPDATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
 RolePolicy.allow_write([
     'name',
     'permission_ids',

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/OauthPresenter.py

@@ -7,4 +7,5 @@ class OauthPresenter(BasePresenter):
         'expires',
         'provider',
         'url',
+        'redirect_query_params',
     ]

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/BackfillResource.py

@@ -16,6 +16,7 @@ ALLOWED_PAYLOAD_KEYS = [
     'interval_units',
     'name',
     'start_datetime',
+    'variables',
 ]
 
 

mage_ai/api/resources/BlockResource.py

@@ -22,6 +22,7 @@ from mage_ai.data_preparation.models.constants import (
     FILE_EXTENSION_TO_BLOCK_LANGUAGE,
     BlockLanguage,
     BlockType,
+    PipelineType,
 )
 from mage_ai.data_preparation.models.custom_templates.custom_block_template import (
     CustomBlockTemplate,
@@ -62,6 +63,62 @@ class BlockResource(GenericResource):
             block_mapping = {}
 
             pipeline = parent_model.pipeline
+            is_data_integration_pipeline = pipeline and PipelineType.INTEGRATION == pipeline.type
+
+            if is_data_integration_pipeline:
+                original_blocks_mapping = {}
+
+                for block_run in parent_model.block_runs:
+                    block_run_block_uuid = block_run.block_uuid
+                    block = pipeline.get_block(block_run_block_uuid)
+                    if not block:
+                        continue
+
+                    if block.uuid not in original_blocks_mapping:
+                        original_blocks_mapping[block.uuid] = {}
+
+                    block_dict = block.to_dict()
+                    block_dicts_by_uuid[block_run_block_uuid] = block_dict
+                    original_blocks_mapping[block.uuid][block_run_block_uuid] = block_dict
+
+                for block_run_block_uuid, block_dict in block_dicts_by_uuid.items():
+                    block_uuid = block_dict.get('uuid')
+                    block_uuid_parts = block_run_block_uuid.split(':')
+                    group_parts = [part for part in block_uuid_parts if part != block_uuid]
+                    group_uuid = ':'.join(group_parts)
+
+                    for key in [
+                        'downstream_blocks',
+                        'upstream_blocks',
+                    ]:
+                        uuids = []
+                        for block_uuid2 in block_dict.get(key):
+                            mapping = original_blocks_mapping.get(block_uuid2)
+                            if mapping:
+                                for uuid2 in list(mapping.keys()):
+                                    if uuid2 == ':'.join([block_uuid2, group_uuid]):
+                                        uuids.append(uuid2)
+                            else:
+                                uuids.append(block_uuid2)
+
+                        block_dict[key] = uuids
+
+                    if len(group_parts) >= 1:
+                        block_dict['tags'] = [group_parts[0]]
+
+                    if len(group_parts) >= 2:
+                        block_dict['description'] = group_parts[1]
+
+                    block_dict['name'] = block_uuid
+                    block_dict['uuid'] = block_run_block_uuid
+                    block_dicts_by_uuid[block_run_block_uuid] = block_dict
+
+                return self.build_result_set(
+                    block_dicts_by_uuid.values(),
+                    user,
+                    **kwargs,
+                )
+
             for block_run in parent_model.block_runs:
                 block_run_block_uuid = block_run.block_uuid
                 if block_run_block_uuid not in block_mapping:

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(