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

mage_ai/api/policies/BasePolicy.py

@@ -1,15 +1,18 @@
 import importlib
 import inspect
 from collections.abc import Iterable
-from typing import List, Tuple, Union
+from typing import Callable, List, Tuple, Union
 
 import inflection
 
 from mage_ai import settings
 from mage_ai.api.constants import AttributeOperationType, AttributeType
 from mage_ai.api.errors import ApiError
+from mage_ai.api.mixins.result_set import ResultSetMixIn
 from mage_ai.api.oauth_scope import OauthScope
 from mage_ai.api.operations.constants import OperationType
+from mage_ai.api.policies.mixins.user_permissions import UserPermissionMixIn
+from mage_ai.api.result_set import ResultSet
 from mage_ai.api.utils import (
     has_at_least_admin_role,
     has_at_least_editor_role,
@@ -21,15 +24,23 @@ from mage_ai.api.utils import (
 from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.constants import Entity
 from mage_ai.services.tracking.metrics import increment
-from mage_ai.settings import DISABLE_NOTEBOOK_EDIT_ACCESS, REQUIRE_USER_AUTHENTICATION
+from mage_ai.settings import (
+    DISABLE_NOTEBOOK_EDIT_ACCESS,
+    REQUIRE_USER_AUTHENTICATION,
+    REQUIRE_USER_PERMISSIONS,
+)
 from mage_ai.shared.hash import extract
 
 
-class BasePolicy():
+class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
     action_rules = {}
     query_rules = {}
     read_rules = {}
     write_rules = {}
+    override_permission_action_rules = {}
+    override_permission_query_rules = {}
+    override_permission_read_rules = {}
+    override_permission_write_rules = {}
 
     def __init__(self, resource, current_user, **kwargs):
         self.current_user = current_user
@@ -37,10 +48,15 @@ class BasePolicy():
         self.resource = resource
         self.parent_model_attr = None
         self.parent_resource_attr = None
-        if isinstance(resource, Iterable):
-            self.resources = resource
+        self.result_set_attr = None
+
+        if resource:
+            if isinstance(resource, Iterable):
+                self.resources = resource
+            else:
+                self.resources = [resource]
         else:
-            self.resources = [resource]
+            self.result_set_attr = ResultSet([])
 
     @property
     def entity(self) -> Tuple[Union[Entity, None], Union[str, None]]:
@@ -48,60 +64,142 @@ class BasePolicy():
 
     @classmethod
     def action_rule(self, action):
+        if REQUIRE_USER_PERMISSIONS:
+            return self.action_rule_with_permissions(action)
+
         if not self.action_rules.get(self.__name__):
             self.action_rules[self.__name__] = {}
         return self.action_rules[self.__name__].get(action)
 
+    @classmethod
+    def override_permission_action_rule(self, action):
+        if not self.override_permission_action_rules.get(self.__name__):
+            self.override_permission_action_rules[self.__name__] = {}
+        return self.override_permission_action_rules[self.__name__].get(action)
+
     @classmethod
     def query_rule(self, query):
+        if REQUIRE_USER_PERMISSIONS:
+            return self.attribute_rule_with_permissions(AttributeOperationType.QUERY, query)
+
         if not self.query_rules.get(self.__name__):
             self.query_rules[self.__name__] = {}
         return self.query_rules[self.__name__].get(query)
 
+    @classmethod
+    def override_permission_query_rule(self, query):
+        if not self.override_permission_query_rules.get(self.__name__):
+            self.override_permission_query_rules[self.__name__] = {}
+        return self.override_permission_query_rules[self.__name__].get(query)
+
     @classmethod
     def read_rule(self, read):
+        if REQUIRE_USER_PERMISSIONS:
+            return self.attribute_rule_with_permissions(AttributeOperationType.READ, read)
+
         if not self.read_rules.get(self.__name__):
             self.read_rules[self.__name__] = {}
         return self.read_rules[self.__name__].get(read)
 
+    @classmethod
+    def override_permission_read_rule(self, read):
+        if not self.override_permission_read_rules.get(self.__name__):
+            self.override_permission_read_rules[self.__name__] = {}
+        return self.override_permission_read_rules[self.__name__].get(read)
+
     @classmethod
     def write_rule(self, write):
+        if REQUIRE_USER_PERMISSIONS:
+            return self.attribute_rule_with_permissions(AttributeOperationType.WRITE, write)
+
         if not self.write_rules.get(self.__name__):
             self.write_rules[self.__name__] = {}
         return self.write_rules[self.__name__].get(write)
 
+    @classmethod
+    def override_permission_write_rule(self, write):
+        if not self.override_permission_write_rules.get(self.__name__):
+            self.override_permission_write_rules[self.__name__] = {}
+        return self.override_permission_write_rules[self.__name__].get(write)
+
     @classmethod
     def allow_actions(self, array, **kwargs):
         if not self.action_rules.get(self.__name__):
             self.action_rules[self.__name__] = {}
-        for key in array:
+        if not self.override_permission_action_rules.get(self.__name__):
+            self.override_permission_action_rules[self.__name__] = {}
+
+        array_use = array or [OperationType.ALL]
+        for key in array_use:
             if not self.action_rules[self.__name__].get(key):
                 self.action_rules[self.__name__][key] = {}
             for scope in kwargs.get('scopes', []):
                 self.action_rules[self.__name__][key][scope] = extract(kwargs, [
                                                                        'condition'])
+            if not self.override_permission_action_rules[self.__name__].get(key):
+                self.override_permission_action_rules[self.__name__][key] = {}
+            for scope in kwargs.get('scopes', []):
+                self.override_permission_action_rules[self.__name__][key][scope] = extract(
+                    kwargs,
+                    [
+                        'override_permission_condition',
+                    ],
+                )
 
     @classmethod
     def allow_query(self, array: List = None, **kwargs):
         if not self.query_rules.get(self.__name__):
             self.query_rules[self.__name__] = {}
+        if not self.override_permission_query_rules.get(self.__name__):
+            self.override_permission_query_rules[self.__name__] = {}
 
         array_use = array or [AttributeType.ALL]
         for key in array_use:
             if not self.query_rules[self.__name__].get(key):
                 self.query_rules[self.__name__][key] = {}
+            if not self.override_permission_query_rules[self.__name__].get(key):
+                self.override_permission_query_rules[self.__name__][key] = {}
+
+            actions = kwargs.get('on_action', [OperationType.ALL])
+            actions = actions if isinstance(actions, list) else [actions]
+
             for scope in kwargs.get('scopes', []):
-                self.query_rules[self.__name__][key][scope] = extract(kwargs, ['condition'])
+                if not self.query_rules[self.__name__][key].get(scope):
+                    self.query_rules[self.__name__][key][scope] = {}
+                for action in actions:
+                    self.query_rules[self.__name__][key][scope][action] = extract(
+                        kwargs,
+                        [
+                            'condition',
+                        ],
+                    )
+
+                if not self.override_permission_query_rules[self.__name__][key].get(scope):
+                    self.override_permission_query_rules[self.__name__][key][scope] = {}
+                for action in actions:
+                    self.override_permission_query_rules[self.__name__][key][scope][action] = \
+                        extract(
+                            kwargs,
+                            [
+                                'override_permission_condition',
+                            ])
 
     @classmethod
     def allow_read(self, array, **kwargs):
         if not self.read_rules.get(self.__name__):
             self.read_rules[self.__name__] = {}
+        if not self.override_permission_read_rules.get(self.__name__):
+            self.override_permission_read_rules[self.__name__] = {}
+
         for key in array:
             if not self.read_rules[self.__name__].get(key):
                 self.read_rules[self.__name__][key] = {}
+            if not self.override_permission_read_rules[self.__name__].get(key):
+                self.override_permission_read_rules[self.__name__][key] = {}
+
             actions = kwargs.get('on_action', [OperationType.ALL])
             actions = actions if isinstance(actions, list) else [actions]
+
             for scope in kwargs.get('scopes', []):
                 if not self.read_rules[self.__name__][key].get(scope):
                     self.read_rules[self.__name__][key][scope] = {}
@@ -109,15 +207,32 @@ class BasePolicy():
                     self.read_rules[self.__name__][key][scope][action] = extract(kwargs, [
                                                                                  'condition'])
 
+                if not self.override_permission_read_rules[self.__name__][key].get(scope):
+                    self.override_permission_read_rules[self.__name__][key][scope] = {}
+                for action in actions:
+                    self.override_permission_read_rules[self.__name__][key][scope][action] = \
+                        extract(
+                            kwargs,
+                            [
+                                'override_permission_condition',
+                            ])
+
     @classmethod
     def allow_write(self, array, **kwargs):
         if not self.write_rules.get(self.__name__):
             self.write_rules[self.__name__] = {}
+        if not self.override_permission_write_rules.get(self.__name__):
+            self.override_permission_write_rules[self.__name__] = {}
+
         for key in array:
             if not self.write_rules[self.__name__].get(key):
                 self.write_rules[self.__name__][key] = {}
+            if not self.override_permission_write_rules[self.__name__].get(key):
+                self.override_permission_write_rules[self.__name__][key] = {}
+
             actions = kwargs.get('on_action', [OperationType.ALL])
             actions = actions if isinstance(actions, list) else [actions]
+
             for scope in kwargs.get('scopes', []):
                 if not self.write_rules[self.__name__][key].get(scope):
                     self.write_rules[self.__name__][key][scope] = {}
@@ -125,10 +240,24 @@ class BasePolicy():
                     self.write_rules[self.__name__][key][scope][action] = extract(kwargs, [
                                                                                   'condition'])
 
+                if not self.override_permission_write_rules[self.__name__][key].get(scope):
+                    self.override_permission_write_rules[self.__name__][key][scope] = {}
+                for action in actions:
+                    self.override_permission_write_rules[self.__name__][key][scope][action] = \
+                        extract(
+                            kwargs,
+                            [
+                                'override_permission_condition',
+                            ])
+
     @classmethod
     def resource_name(self):
         return inflection.pluralize(self.resource_name_singular())
 
+    @classmethod
+    def model_name(self) -> str:
+        return self.__name__.replace('Policy', '')
+
     @classmethod
     def resource_name_singular(self):
         return inflection.underscore(
@@ -178,6 +307,9 @@ class BasePolicy():
         )
 
     async def authorize_action(self, action):
+        if self.is_owner():
+            return True
+
         config = self.__class__.action_rule(action)
         if config:
             await self.__validate_scopes(action, config.keys())
@@ -187,6 +319,9 @@ class BasePolicy():
                     action,
                     config[self.current_scope()]['condition'],
                     operation=action,
+                    override_permission_condition=(self.__class__.override_permission_action_rule(
+                        action,
+                    ) or {}).get(self.current_scope(), {}).get('override_permission_condition'),
                 )
         else:
             error = ApiError.UNAUTHORIZED_ACCESS
@@ -207,9 +342,11 @@ class BasePolicy():
         if 'read' == read_or_write:
             attribute_operation = AttributeOperationType.READ
             orig_config = self.__class__.read_rule(attrb)
+            orig_config_override = self.__class__.override_permission_read_rule(attrb)
         else:
-            orig_config = self.__class__.write_rule(attrb)
             attribute_operation = AttributeOperationType.WRITE
+            orig_config = self.__class__.write_rule(attrb)
+            orig_config_override = self.__class__.override_permission_write_rule(attrb)
 
         config = None
         if orig_config:
@@ -220,6 +357,15 @@ class BasePolicy():
             if config is None:
                 config = config_scope.get(OperationType.ALL)
 
+        config_override = None
+        if orig_config_override:
+            await self.__validate_scopes(attrb, orig_config_override.keys())
+            config_override_scope = orig_config_override.get(self.current_scope(), {})
+            config_override = config_override_scope.get(api_operation_action)
+
+            if config_override is None:
+                config_override = config_override_scope.get(OperationType.ALL)
+
         if config is None:
             error = ApiError.UNAUTHORIZED_ACCESS
             error.update({
@@ -231,13 +377,16 @@ class BasePolicy():
                 ),
             })
             raise ApiError(error)
+
         cond = config.get('condition')
+
         if cond:
             await self.__validate_condition(
                 attrb,
                 cond,
                 attribute_operation=attribute_operation,
                 operation=api_operation_action,
+                override_permission_condition=config_override.get('override_permission_condition'),
                 **kwargs,
             )
 
@@ -248,36 +397,63 @@ class BasePolicy():
         for attrb in attrbs:
             await self.authorize_attribute(read_or_write, attrb, **kwargs)
 
-    async def authorize_query(self, query):
-        if not query:
-            return
+    async def authorize_query(self, query, **kwargs):
+        if self.is_owner() or not query:
+            return True
+
+        api_operation_action = self.options.get(
+            'api_operation_action',
+            kwargs.get('api_operation_action', OperationType.ALL),
+        )
 
         for key, value in query.items():
-            if key != settings.QUERY_API_KEY:
-                api_operation_action = self.options.get(
-                    'api_operation_action',
-                    OperationType.ALL,
-                )
+            if key == settings.QUERY_API_KEY:
+                continue
 
-                error_message = f'Query parameter {key} of value {value} ' \
-                    f'is not permitted on {api_operation_action} operation.'
-
-                config = self.__class__.query_rule(key) or \
-                    self.__class__.query_rule(AttributeType.ALL)
-
-                if not config:
-                    error = ApiError.UNAUTHORIZED_ACCESS
-                    error.update({
-                        'message': error_message,
-                    })
-                    raise ApiError(error)
-                elif config.get(self.current_scope(), {}).get('condition'):
-                    await self.__validate_condition(
-                        key,
-                        config[self.current_scope()]['condition'],
-                        message=error_message,
-                        operation=api_operation_action,
-                    )
+            error_message = f'Query parameter {key} of value {value} ' \
+                f'is not permitted on {api_operation_action} operation ' \
+                f'for {self.__class__.resource_name()}.'
+
+            orig_config = self.__class__.query_rule(key) or \
+                self.__class__.query_rule(AttributeType.ALL)
+
+            config = None
+            if orig_config:
+                await self.__validate_scopes(key, orig_config.keys())
+                config_scope = orig_config.get(self.current_scope(), {})
+                config = config_scope.get(api_operation_action)
+
+                if config is None:
+                    config = config_scope.get(OperationType.ALL)
+
+            orig_config_override = self.__class__.override_permission_query_rule(key) or \
+                self.__class__.override_permission_query_rule(AttributeType.ALL)
+
+            config_override = None
+            if orig_config_override:
+                await self.__validate_scopes(key, orig_config_override.keys())
+                config_override_scope = orig_config_override.get(self.current_scope(), {})
+                config_override = config_override_scope.get(api_operation_action)
+
+                if config_override is None:
+                    config_override = config_override_scope.get(OperationType.ALL)
+
+            if config is None:
+                error = ApiError.UNAUTHORIZED_ACCESS
+                error.update({
+                    'message': error_message,
+                })
+                raise ApiError(error)
+
+            cond = config.get('condition')
+
+            await self.__validate_condition(
+                key,
+                cond,
+                message=error_message,
+                operation=api_operation_action,
+                override_permission_condition=config_override.get('override_permission_condition'),
+            )
 
     def parent_model(self):
         if not self.parent_model_attr:
@@ -306,12 +482,19 @@ class BasePolicy():
         else:
             return OauthScope.CLIENT_PUBLIC
 
+    def result_set(self) -> ResultSet:
+        if self.resource:
+            return self.resource.result_set()
+
+        return self.result_set_attr
+
     async def __validate_condition(
         self,
         action,
         cond,
         attribute_operation: AttributeOperationType = None,
         operation: OperationType = None,
+        override_permission_condition: Callable = None,
         **kwargs,
     ):
         if not cond:
@@ -321,16 +504,22 @@ class BasePolicy():
         if validation and inspect.isawaitable(validation):
             validation = await validation
 
+        if not validation and override_permission_condition:
+            validation = override_permission_condition(self)
+            if validation and inspect.isawaitable(validation):
+                validation = await validation
+
         if not validation:
+            r_name = self.resource_name()
             error = ApiError.UNAUTHORIZED_ACCESS
-            message = f'Unauthorized access for {action}'
+            message = f'Unauthorized access for {action} on {r_name}'
 
             if attribute_operation:
-                message = f'Unauthorized {attribute_operation} access for {action}'
+                message = f'Unauthorized {attribute_operation} access for {action} on {r_name}'
                 if operation:
-                    message = f'{message} on {operation} operation'
+                    message = f'{message} for {operation} operation'
             elif operation:
-                message = f'Unauthorized access for {action} on {operation} operation'
+                message = f'Unauthorized operation {operation} on {r_name}'
 
             error.update({
                 'message': kwargs.get(

mage_ai/api/policies/PermissionPolicy.py

@@ -26,6 +26,8 @@ PermissionPolicy.allow_actions([
 
 
 PermissionPolicy.allow_read(PermissionPresenter.default_attributes + [
+    'entity_names',
+    'entity_types',
     'query_attributes',
     'read_attributes',
     'write_attributes',
@@ -41,11 +43,16 @@ PermissionPolicy.allow_read(PermissionPresenter.default_attributes + [
 
 PermissionPolicy.allow_read(PermissionPresenter.default_attributes + [
     'role',
+    'user',
     'user_id',
+    'users',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
+    constants.CREATE,
+    constants.DELETE,
     constants.DETAIL,
+    constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
 
@@ -56,9 +63,19 @@ PermissionPolicy.allow_write([
     'entity_name',
     'entity_type',
     'role_id',
+    'role_ids',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.CREATE,
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_admin_role())
+
+
+PermissionPolicy.allow_query([
+    'only_entity_options',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/PipelineSchedulePolicy.py

@@ -1,8 +1,10 @@
 from mage_ai.api.oauth_scope import OauthScope
 from mage_ai.api.operations import constants
 from mage_ai.api.policies.BasePolicy import BasePolicy
-from mage_ai.api.policies.utils import validate_pipeline_interactions_permissions
 from mage_ai.api.presenters.PipelineSchedulePresenter import PipelineSchedulePresenter
+from mage_ai.data_preparation.models.pipelines.interactions import PipelineInteractions
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.constants import Entity
 
@@ -40,7 +42,21 @@ async def authorize_operation_create(policy: PipelineSchedulePolicy) -> bool:
     if policy.has_at_least_editor_role():
         return True
 
-    cond = await validate_pipeline_interactions_permissions(policy)
+    pipeline = policy.parent_model()
+    if not pipeline:
+        pipeline = policy.resource.pipeline
+
+    if not Project(pipeline.repo_config).is_feature_enabled(FeatureUUID.INTERACTIONS):
+        return False
+
+    pipeline_interaction = PipelineInteractions(pipeline)
+    cond = await pipeline_interaction.filter_for_permissions(policy.current_user)
+
+    if not policy.result_set().context.data:
+        policy.result_set().context.data = {}
+    if not policy.result_set().context.data.get('pipeline_interactions'):
+        policy.result_set().context.data['pipeline_interactions'] = {}
+    policy.result_set().context.data['pipeline_interactions'][pipeline.uuid] = pipeline_interaction
 
     return cond
 

mage_ai/api/policies/RolePermissionPolicy.py

@@ -26,8 +26,9 @@ RolePermissionPolicy.allow_read(RolePermissionPresenter.default_attributes + [
 
 
 RolePermissionPolicy.allow_write([
+    'permission_id',
+    'permission_ids',
     'role_id',
-    'user_id',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[

mage_ai/api/policies/RolePolicy.py

@@ -44,6 +44,7 @@ RolePolicy.allow_read([
 
 RolePolicy.allow_read(RolePresenter.default_attributes + [
     'user_id',
+    'users',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
@@ -55,6 +56,8 @@ RolePolicy.allow_read(RolePresenter.default_attributes + [
 
 RolePolicy.allow_write([
     'name'
+    'permission_ids',
+    'user_ids',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[

mage_ai/api/policies/SessionPolicy.py

@@ -1,10 +1,45 @@
+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.orchestration.constants import Entity
+from mage_ai.shared.hash import merge_dict
+
+
+class SessionPolicy(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: OperationType.CREATE == operation,
+            ),
+        })
+
+    @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_PUBLIC: {
+                OperationType.CREATE: config[OauthScope.CLIENT_PUBLIC][OperationType.CREATE],
+            },
+        })
 
-class SessionPolicy(BasePolicy):
     @property
     def entity(self):
         return Entity.ANY, None