mage-ai 0.9.32__py3-none-any.whl → 0.9.34__py3-none-any.whl

mage_ai/api/policies/StatusPolicy.py

@@ -1,11 +1,50 @@
+from typing import Dict
+
+from mage_ai.api.constants import AttributeOperationType
 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.StatusPresenter import StatusPresenter
+from mage_ai.shared.hash import merge_dict
+
 
+class StatusPolicy(BasePolicy, UserPermissionMixIn):
+    @classmethod
+    def action_rule_with_permissions(self, operation: OperationType) -> Dict:
+        return merge_dict(super().action_rule_with_permissions(operation), {
+            OauthScope.CLIENT_PRIVATE: dict(
+                condition=lambda _policy: OperationType.LIST == operation,
+            ),
+            OauthScope.CLIENT_PUBLIC: dict(
+                condition=lambda _policy: OperationType.LIST == operation,
+            ),
+        })
 
-class StatusPolicy(BasePolicy):
-    pass
+    @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)
+        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_PRIVATE: {
+                OperationType.LIST: config[OauthScope.CLIENT_PRIVATE][OperationType.LIST],
+            },
+            OauthScope.CLIENT_PUBLIC: {
+                OperationType.LIST: config[OauthScope.CLIENT_PUBLIC][OperationType.LIST],
+            },
+        })
 
 
 StatusPolicy.allow_actions([
@@ -15,7 +54,11 @@ StatusPolicy.allow_actions([
     OauthScope.CLIENT_PUBLIC,
 ])
 
-StatusPolicy.allow_read(StatusPresenter.default_attributes, scopes=[
+StatusPolicy.allow_read(StatusPresenter.default_attributes + [
+    'active_pipeline_run_count',
+    'last_scheduler_activity',
+    'last_user_request',
+], scopes=[
     OauthScope.CLIENT_PRIVATE,
     OauthScope.CLIENT_PUBLIC,
 ], on_action=[

mage_ai/api/policies/UserPolicy.py

@@ -16,6 +16,7 @@ UserPolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.is_owner())
 
+
 UserPolicy.allow_actions([
     constants.DETAIL,
     constants.UPDATE,
@@ -23,24 +24,40 @@ UserPolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.is_current_user() or policy.has_at_least_admin_role())
 
+
 UserPolicy.allow_actions([
     constants.LIST,
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_admin_role())
 
+
 UserPolicy.allow_read(UserPresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.is_current_user() or policy.has_at_least_admin_role())
 
+
 UserPolicy.allow_read(UserPresenter.default_attributes + [
+    'permissions',
+    'token',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.DETAIL,
+    constants.UPDATE,
+], condition=lambda policy: policy.is_current_user() or policy.has_at_least_admin_role())
+
+
+UserPolicy.allow_read(UserPresenter.default_attributes + [
+    'permissions',
     'token',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.CREATE,
     constants.DELETE,
-], condition=lambda policy: policy.is_current_user() or policy.is_owner())
+], condition=lambda policy: policy.is_owner())
+
 
 UserPolicy.allow_write([
     'avatar',
@@ -57,15 +74,19 @@ UserPolicy.allow_write([
     constants.UPDATE,
 ], condition=lambda policy: policy.is_current_user() or policy.has_at_least_admin_role())
 
+
 UserPolicy.allow_write([
+    'role_ids',
     'roles',
-    'roles_new'
+    'roles_new',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
+    constants.DELETE,
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_admin_role())
 
+
 UserPolicy.allow_write([
     'owner',
 ], scopes=[
@@ -74,6 +95,7 @@ UserPolicy.allow_write([
     constants.UPDATE,
 ], condition=lambda policy: policy.is_owner())
 
+
 UserPolicy.allow_write([
     'avatar',
     'email',
@@ -81,6 +103,7 @@ UserPolicy.allow_write([
     'last_name',
     'password',
     'password_confirmation',
+    'role_ids',
     'roles',
     'roles_new',
     'username',

mage_ai/api/policies/mixins/user_permissions.py

@@ -0,0 +1,241 @@
+import asyncio
+from typing import Any, Callable, Dict, Tuple
+
+from mage_ai.api.constants import (
+    ATTRIBUTE_OPERATION_TYPE_DISABLE_TO_ACCESS_MAPPING,
+    ATTRIBUTE_OPERATION_TYPE_TO_ACCESS_MAPPING,
+    OPERATION_TYPE_DISABLE_TO_ACCESS_MAPPING,
+    OPERATION_TYPE_TO_ACCESS_MAPPING,
+    AttributeOperationType,
+)
+from mage_ai.api.oauth_scope import OauthScope
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.authentication.permissions.constants import (
+    PERMISSION_ACCESS_WITH_MULTIPLE_ACCESS,
+    RESERVED_ENTITY_NAMES,
+    EntityName,
+    PermissionAccess,
+)
+from mage_ai.orchestration.db.models.oauth import Permission
+
+
+def get_entity_name_from_policy(policy) -> EntityName:
+    model_name = policy.model_name()
+    if model_name in EntityName._value2member_map_:
+        return EntityName(model_name)
+
+    return None
+
+
+async def validate_condition_with_permissions(
+    policy,
+    operation: OperationType,
+    attribute_operation_type: AttributeOperationType = None,
+    resource_attribute: str = None,
+) -> bool:
+    entity_name = get_entity_name_from_policy(policy)
+    access = OPERATION_TYPE_TO_ACCESS_MAPPING.get(operation) or Permission.Access.OWNER
+    disable_access = OPERATION_TYPE_DISABLE_TO_ACCESS_MAPPING.get(operation)
+
+    access_for_attribute_operation = None
+    disable_access_for_attribute_operation = None
+    if attribute_operation_type:
+        access_for_attribute_operation = ATTRIBUTE_OPERATION_TYPE_TO_ACCESS_MAPPING.get(
+            attribute_operation_type,
+        )
+        disable_access_for_attribute_operation = \
+            ATTRIBUTE_OPERATION_TYPE_DISABLE_TO_ACCESS_MAPPING.get(
+                attribute_operation_type,
+            )
+
+    permissions = await (policy.resource or policy).load_and_cache_user_permissions()
+
+    async def __permission_grants_access(
+        permission: Permission,
+        entity_name=entity_name,
+        access=access,
+        disable_access=disable_access,
+        access_for_attribute_operation=access_for_attribute_operation,
+        disable_access_for_attribute_operation=disable_access_for_attribute_operation,
+        attribute_operation_type=attribute_operation_type,
+        resource_attribute=resource_attribute,
+    ) -> Tuple[bool, bool]:
+        if permission.access is None:
+            return (False, False)
+
+        # Check if user is an owner
+        if permission.access & PermissionAccess.OWNER:
+            return (True, False)
+
+        # 1. Get permissions for current entity_name for roles belonging to current user.
+        # Include permissions where entity_name is ALL or ALL_EXCEPT_RESERVED.
+        correct_entity_name = permission.entity_name == entity_name or \
+            permission.entity_name == EntityName.ALL or \
+            (
+                permission.entity_name == EntityName.ALL_EXCEPT_RESERVED and
+                entity_name not in RESERVED_ENTITY_NAMES
+            )
+
+        if not correct_entity_name:
+            return (False, False)
+
+        # 2a. Don’t grant access if permission disables access to this entity for this operation.
+        if disable_access is not None and permission.access & disable_access:
+            return (False, True)
+
+        # 3. Add additional permission access (e.g. read, list, detail for viewer)
+        # to grant access to this entity and its attributes
+        permission_accesses = PERMISSION_ACCESS_WITH_MULTIPLE_ACCESS.get(f'{permission.access}')
+        if not permission_accesses:
+            permission_accesses = [permission.access]
+
+        permission_granted = False
+        permission_disabled = False
+
+        for permission_access in permission_accesses:
+            if permission_granted:
+                break
+
+            # Access to all operations and attribute operations
+            if permission_access & PermissionAccess.ALL.value:
+                permission_granted = True
+                break
+
+            # Disable all operations
+            if permission_access & PermissionAccess.DISABLE_OPERATION_ALL.value:
+                permission_disabled = True
+                break
+
+            has_access_for_all_operations = permission_access & PermissionAccess.OPERATION_ALL
+            valid_for_operation = has_access_for_all_operations or permission_access & access
+
+            # If this condition is validating attribute operations for an attribute:
+            if access_for_attribute_operation:
+                access_for_all = 0
+                disable_access_for_all = 0
+
+                if AttributeOperationType.QUERY == attribute_operation_type:
+                    access_for_all = PermissionAccess.QUERY_ALL
+                    disable_access_for_all = PermissionAccess.DISABLE_QUERY_ALL
+                elif AttributeOperationType.READ == attribute_operation_type:
+                    access_for_all = PermissionAccess.READ_ALL
+                    disable_access_for_all = PermissionAccess.DISABLE_READ_ALL
+                elif AttributeOperationType.WRITE == attribute_operation_type:
+                    access_for_all = PermissionAccess.WRITE_ALL
+                    disable_access_for_all = PermissionAccess.DISABLE_WRITE_ALL
+
+                if permission_access & disable_access_for_all:
+                    permission_disabled = True
+                    break
+
+                has_access_for_all_attributes = permission_access & access_for_all
+                valid_for_operation = valid_for_operation and \
+                    (
+                        permission_access & access_for_attribute_operation or
+                        has_access_for_all_attributes
+                    )
+
+                # Don’t grant access if permission disables access to this entity’s attributes
+                # for this attribute operation.
+                if disable_access_for_attribute_operation is not None and \
+                        permission.access & disable_access_for_attribute_operation:
+
+                    disabled_attributes = []
+                    if AttributeOperationType.QUERY == attribute_operation_type:
+                        disabled_attributes = permission.query_attributes
+                    elif AttributeOperationType.READ == attribute_operation_type:
+                        disabled_attributes = permission.read_attributes
+                    elif AttributeOperationType.WRITE == attribute_operation_type:
+                        disabled_attributes = permission.write_attributes
+
+                    if resource_attribute in (disabled_attributes or []):
+                        permission_disabled = True
+                        break
+
+                if not has_access_for_all_attributes:
+                    if valid_for_operation and attribute_operation_type and resource_attribute:
+                        permitted_attributes = []
+                        if AttributeOperationType.QUERY == attribute_operation_type:
+                            permitted_attributes = permission.query_attributes
+                        elif AttributeOperationType.READ == attribute_operation_type:
+                            permitted_attributes = permission.read_attributes
+                        elif AttributeOperationType.WRITE == attribute_operation_type:
+                            permitted_attributes = permission.write_attributes
+
+                        valid_for_operation = resource_attribute in (permitted_attributes or [])
+
+            if valid_for_operation:
+                permission_granted = True
+                break
+
+        return (permission_granted, permission_disabled)
+
+    validation_results = await asyncio.gather(
+        *[__permission_grants_access(p) for p in permissions]
+    )
+
+    authorized = False
+    unauthorized = False
+    for permission_granted, permission_disabled in validation_results:
+        if permission_disabled:
+            unauthorized = True
+            break
+
+        if permission_granted and not authorized:
+            authorized = True
+
+    return authorized and not unauthorized
+
+
+class UserPermissionMixIn:
+    @classmethod
+    def action_rule_with_permissions(self, operation: OperationType) -> Dict:
+        return {
+            OauthScope.CLIENT_PRIVATE: dict(
+                condition=self.build_validate_condition(operation),
+            ),
+        }
+
+    @classmethod
+    def attribute_rule_with_permissions(
+        self,
+        attribute_operation_type: AttributeOperationType,
+        resource_attribute: str,
+    ) -> Dict:
+        conditions = {}
+        for operation in OperationType:
+            conditions[operation.value] = dict(
+                condition=self.build_validate_attribute(
+                    operation,
+                    attribute_operation_type,
+                    resource_attribute,
+                ),
+            )
+
+        return {
+            OauthScope.CLIENT_PRIVATE: conditions,
+        }
+
+    @classmethod
+    def build_validate_attribute(
+        self,
+        operation: OperationType,
+        attribute_operation_type: AttributeOperationType,
+        resource_attribute: str,
+    ) -> Callable[[Any], bool]:
+        def _validate_condition(policy) -> bool:
+            return validate_condition_with_permissions(
+                policy,
+                operation,
+                attribute_operation_type=attribute_operation_type,
+                resource_attribute=resource_attribute,
+            )
+
+        return _validate_condition
+
+    @classmethod
+    def build_validate_condition(self, operation: OperationType) -> Callable[[Any], bool]:
+        def _validate_condition(policy) -> bool:
+            return validate_condition_with_permissions(policy, operation)
+
+        return _validate_condition

mage_ai/api/presenters/PermissionPresenter.py

@@ -1,5 +1,12 @@
+from typing import List
+
 from mage_ai.api.operations import constants
 from mage_ai.api.presenters.BasePresenter import BasePresenter
+from mage_ai.authentication.permissions.constants import (
+    BlockEntityType,
+    EntityName,
+    PipelineEntityType,
+)
 
 
 class PermissionPresenter(BasePresenter):
@@ -14,6 +21,14 @@ class PermissionPresenter(BasePresenter):
         'updated_at',
     ]
 
+    def entity_names(self, **kwargs) -> List[str]:
+        return sorted([n for n in EntityName])
+
+    def entity_types(self, **kwargs) -> List[str]:
+        return sorted(
+            [n for n in BlockEntityType] + [n for n in PipelineEntityType],
+        )
+
 
 PermissionPresenter.register_format(
     f'role/{constants.DETAIL}',
@@ -27,14 +42,26 @@ PermissionPresenter.register_format(
 
 PermissionPresenter.register_formats([
     constants.CREATE,
+    constants.DELETE,
     constants.DETAIL,
     constants.UPDATE,
 ], PermissionPresenter.default_attributes + [
-        'query_attributes',
-        'read_attributes',
-        'role',
-        'roles',
-        'user_id',
-        'write_attributes',
-    ],
-)
+    'entity_names',
+    'entity_types',
+    'query_attributes',
+    'read_attributes',
+    'role',
+    'roles',
+    'user',
+    'user_id',
+    'users',
+    'write_attributes',
+])
+
+
+PermissionPresenter.register_formats([
+    'with_only_entity_options',
+], [
+    'entity_names',
+    'entity_types',
+])

mage_ai/api/presenters/RolePresenter.py

@@ -10,6 +10,7 @@ class RolePresenter(BasePresenter):
         'permissions',
         'role_permissions',
         'updated_at',
+        'user',
     ]
 
 
@@ -25,12 +26,14 @@ RolePresenter.register_format(
 )
 
 
-RolePresenter.register_format(
+RolePresenter.register_formats([
     constants.DETAIL,
-    RolePresenter.default_attributes + [
-        'users',
-    ],
-)
+    constants.CREATE,
+    constants.UPDATE,
+    constants.DELETE,
+], RolePresenter.default_attributes + [
+    'users',
+])
 
 
 RolePresenter.register_formats([
@@ -38,10 +41,24 @@ RolePresenter.register_formats([
     f'permission/{constants.DETAIL}',
     f'permission/{constants.UPDATE}',
 ], [
-        'created_at',
-        'id',
-        'name',
-        'updated_at',
-        'users',
-    ],
-)
+    'created_at',
+    'id',
+    'name',
+    'updated_at',
+    'users',
+])
+
+
+RolePresenter.register_formats([
+    f'user/{constants.CREATE}',
+    f'user/{constants.DELETE}',
+    f'user/{constants.DETAIL}',
+    f'user/{constants.LIST}',
+    f'user/{constants.UPDATE}',
+], [
+    'created_at',
+    'id',
+    'name',
+    'permissions',
+    'updated_at',
+])

mage_ai/api/presenters/StatusPresenter.py

@@ -10,6 +10,17 @@ class StatusPresenter(BasePresenter):
         'instance_type',
         'disable_pipeline_edit_access',
         'require_user_authentication',
+        'require_user_permissions',
         'project_type',
         'project_uuid',
     ]
+
+
+StatusPresenter.register_format(
+    'with_activity_details',
+    StatusPresenter.default_attributes + [
+        'active_pipeline_run_count',
+        'last_scheduler_activity',
+        'last_user_request',
+    ],
+)

mage_ai/api/presenters/UserPresenter.py

@@ -1,46 +1,43 @@
 from mage_ai.api.operations import constants
 from mage_ai.api.presenters.BasePresenter import BasePresenter
-from mage_ai.shared.hash import extract
 
 
 class UserPresenter(BasePresenter):
     default_attributes = [
         'avatar',
+        'created_at',
         'email',
         'first_name',
         'id',
         'last_name',
         'owner',
+        'project_access',
         'roles',
         'roles_display',
         'roles_new',
-        'project_access',
+        'updated_at',
         'username',
     ]
 
-    async def prepare_present(self, **kwargs):
-        data = self.model.to_dict(include_attributes=self.default_attributes)
-        data = extract(data, self.default_attributes, include_blank_values=True)
-
-        roles_new = self.model.roles_new
-        data['roles_new'] = [
-            role.to_dict(include_attributes=['permissions'])
-            for role in roles_new
-        ]
-
-        return data
-
 
-UserPresenter.register_format(
-    constants.CREATE, UserPresenter.default_attributes + ['token', ])
+UserPresenter.register_formats([
+    constants.CREATE,
+    constants.DELETE,
+    constants.DETAIL,
+    constants.UPDATE,
+], UserPresenter.default_attributes + [
+    'permissions',
+    'token',
+])
 
 
-UserPresenter.register_format(
+UserPresenter.register_formats([
+    f'permission/{constants.DETAIL}',
     f'role/{constants.DETAIL}',
-    [
-        'first_name',
-        'id',
-        'last_name',
-        'username',
-    ],
-)
+    f'role/{constants.LIST}',
+], [
+    'first_name',
+    'id',
+    'last_name',
+    'username',
+])

mage_ai/api/resources/BaseResource.py

@@ -5,6 +5,7 @@ import inflection
 
 from mage_ai import settings
 from mage_ai.api.errors import ApiError
+from mage_ai.api.mixins.result_set import ResultSetMixIn
 from mage_ai.api.resources.Resource import Resource
 from mage_ai.api.resources.shared import collective_loaders
 from mage_ai.api.result_set import ResultSet
@@ -12,7 +13,7 @@ from mage_ai.orchestration.db.errors import DoesNotExistError
 from mage_ai.shared.hash import merge_dict
 
 
-class BaseResource(Resource):
+class BaseResource(Resource, ResultSetMixIn):
     collective_loader_attr = {}
     datetime_keys = []
 
@@ -25,11 +26,15 @@ class BaseResource(Resource):
     parent_models_attr = {}
     parent_resource_attr = {}
 
+    @classmethod
+    def model_name(self) -> str:
+        return self.__name__.replace('Resource', '')
+
     @classmethod
     def parser_class(self):
-        klass = self.__name__.replace('Resource', '')
-        module_name = f'mage_ai.api.parsers.{klass}Parser'
-        class_name = f'{klass}Parser'
+        model_name = self.model_name()
+        module_name = f'mage_ai.api.parsers.{model_name}Parser'
+        class_name = f'{model_name}Parser'
 
         try:
             return getattr(importlib.import_module(module_name), class_name)
@@ -38,7 +43,7 @@ class BaseResource(Resource):
 
     @classmethod
     def policy_class(self):
-        model_name = self.__name__.replace('Resource', '')
+        model_name = self.model_name()
         return getattr(
             importlib.import_module(
                 'mage_ai.api.policies.{}Policy'.format(model_name)),
@@ -47,7 +52,7 @@ class BaseResource(Resource):
 
     @classmethod
     def presenter_class(self):
-        model_name = self.__name__.replace('Resource', '')
+        model_name = self.model_name()
         return getattr(
             importlib.import_module(
                 'mage_ai.api.presenters.{}Presenter'.format(model_name)),
@@ -151,7 +156,12 @@ class BaseResource(Resource):
         pass
 
     @classmethod
-    async def process_create(self, payload, user, **kwargs):
+    async def process_create(
+        self,
+        payload,
+        user,
+        **kwargs,
+    ):
         self.on_create_callback = None
         self.on_create_failure_callback = None
         before_create = self.before_create(payload, user, **kwargs)
@@ -259,11 +269,16 @@ class BaseResource(Resource):
 
             raise err
 
-    def result_set(self):
+    def result_set(self) -> ResultSet:
         if self.__result_sets().get(self.__class__.__name__, None):
             return self.__result_sets()[self.__class__.__name__]
         elif not self.result_set_attr:
-            self.result_set_attr = ResultSet([self])
+            if self.result_set_from_external is not None:
+                self.result_set_attr = self.result_set_from_external
+                if self not in self.result_set_attr:
+                    self.result_set_attr.add_results([self])
+            else:
+                self.result_set_attr = ResultSet([self])
         return self.result_set_attr
 
     def update(self, payload, **kwargs):