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

mage_ai/api/mixins/result_set.py

@@ -0,0 +1,65 @@
+from typing import List
+
+from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.oauth import (
+    Permission,
+    Role,
+    RolePermission,
+    UserRole,
+)
+
+CONTEXT_DATA_KEY_USER_PERMISSIONS = '__user_permissions'
+
+
+class ResultSetMixIn:
+    @safe_db_query
+    async def load_and_cache_user_permissions(self) -> List[Permission]:
+        # This will fetch the user permissions and store it on the context data
+        # so that repeat policy user permission validations won’t keep querying the
+        # database for permissions.
+        if not self.current_user:
+            return
+
+        permissions = self.result_set().context.data.get(CONTEXT_DATA_KEY_USER_PERMISSIONS)
+        if permissions:
+            return permissions
+
+        query = (
+            Permission.
+            select(
+                Permission.access,
+                Permission.entity,
+                Permission.entity_id,
+                Permission.entity_name,
+                Permission.entity_type,
+                Permission.id,
+                Permission.options,
+            ).
+            join(
+                RolePermission,
+                RolePermission.permission_id == Permission.id).
+            join(
+                Role,
+                Role.id == RolePermission.role_id).
+            join(
+                UserRole,
+                UserRole.role_id == Role.id).
+            filter(UserRole.user_id == self.current_user.id)
+        )
+
+        permissions = []
+
+        for row in query.all():
+            permission = Permission()
+            permission.access = row.access
+            permission.entity = row.entity
+            permission.entity_id = row.entity_id
+            permission.entity_name = row.entity_name
+            permission.entity_type = row.entity_type
+            permission.id = row.id
+            permission.options = row.options
+            permissions.append(permission)
+
+        self.result_set().context.data[CONTEXT_DATA_KEY_USER_PERMISSIONS] = permissions
+
+        return permissions

mage_ai/api/operations/base.py

@@ -27,6 +27,7 @@ from mage_ai.api.parsers.BaseParser import BaseParser
 from mage_ai.api.presenters.BasePresenter import CustomDict, CustomList
 from mage_ai.api.result_set import ResultSet
 from mage_ai.orchestration.db.errors import DoesNotExistError
+from mage_ai.settings import REQUIRE_USER_PERMISSIONS
 from mage_ai.shared.array import flatten
 from mage_ai.shared.hash import ignore_keys, merge_dict
 from mage_ai.shared.strings import classify
@@ -170,9 +171,13 @@ class BaseOperation():
                 }
         except ApiError as err:
             if err.code == 403 and \
-                    self.user and self.user.project_access == 0 and not self.user.roles:
-                err.message = 'You do not have access to this project. ' + \
-                    'Please ask an admin or owner for permissions.'
+                    self.user and \
+                    self.user.project_access == 0 and \
+                    not self.user.roles:
+
+                if not REQUIRE_USER_PERMISSIONS:
+                    err.message = 'You don’t have access to this project. ' + \
+                        'Please ask an admin or owner for permissions.'
             if settings.DEBUG:
                 raise err
             else:

mage_ai/api/policies/BasePolicy.py

@@ -1,7 +1,7 @@
 import importlib
 import inspect
 from collections.abc import Iterable
-from typing import List, Tuple, Union
+from typing import Callable, List, Tuple, Union
 
 import inflection
 
@@ -37,6 +37,10 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
     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
@@ -67,6 +71,12 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
             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:
@@ -76,6 +86,12 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
             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:
@@ -85,6 +101,12 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
             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:
@@ -94,10 +116,18 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
             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__] = {}
+        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:
@@ -106,18 +136,33 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
             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', []):
                 if not self.query_rules[self.__name__][key].get(scope):
                     self.query_rules[self.__name__][key][scope] = {}
@@ -129,15 +174,32 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                         ],
                     )
 
+                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] = {}
@@ -145,15 +207,32 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                     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] = {}
@@ -161,6 +240,16 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                     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())
@@ -230,6 +319,9 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                     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
@@ -250,9 +342,11 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
         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:
@@ -263,6 +357,15 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
             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({
@@ -274,13 +377,16 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                 ),
             })
             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,
             )
 
@@ -320,6 +426,18 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                 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({
@@ -334,6 +452,7 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
                 cond,
                 message=error_message,
                 operation=api_operation_action,
+                override_permission_condition=config_override.get('override_permission_condition'),
             )
 
     def parent_model(self):
@@ -375,6 +494,7 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
         cond,
         attribute_operation: AttributeOperationType = None,
         operation: OperationType = None,
+        override_permission_condition: Callable = None,
         **kwargs,
     ):
         if not cond:
@@ -384,6 +504,11 @@ class BasePolicy(UserPermissionMixIn, ResultSetMixIn):
         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

mage_ai/api/policies/PermissionPolicy.py

@@ -45,10 +45,14 @@ 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())
 
 
@@ -59,6 +63,7 @@ PermissionPolicy.allow_write([
     'entity_name',
     'entity_type',
     'role_id',
+    'role_ids',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], 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

@@ -135,15 +135,6 @@ async def validate_condition_with_permissions(
                         has_access_for_all_attributes
                     )
 
-                print(
-                    'WTFFFFFFFFFFFFFFFFFFFFFF',
-                    disable_access_for_attribute_operation,
-                    resource_attribute,
-                    permission.id,
-                    permission.access,
-                    permission.access & disable_access_for_attribute_operation,
-                )
-
                 # 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 \

mage_ai/api/presenters/PermissionPresenter.py

@@ -42,6 +42,7 @@ PermissionPresenter.register_format(
 
 PermissionPresenter.register_formats([
     constants.CREATE,
+    constants.DELETE,
     constants.DETAIL,
     constants.UPDATE,
 ], PermissionPresenter.default_attributes + [
@@ -53,6 +54,7 @@ PermissionPresenter.register_formats([
     'roles',
     'user',
     'user_id',
+    'users',
     'write_attributes',
 ])
 

mage_ai/api/presenters/RolePresenter.py

@@ -47,3 +47,18 @@ RolePresenter.register_formats([
     '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/UserPresenter.py

@@ -1,6 +1,5 @@
 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):
@@ -20,21 +19,16 @@ class UserPresenter(BasePresenter):
         '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_formats([

mage_ai/api/resources/PermissionResource.py

@@ -1,7 +1,7 @@
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
-from mage_ai.orchestration.db import safe_db_query
-from mage_ai.orchestration.db.models.oauth import Permission, Role
-from mage_ai.shared.hash import merge_dict
+from mage_ai.orchestration.db import db_connection, safe_db_query
+from mage_ai.orchestration.db.models.oauth import Permission, Role, RolePermission
+from mage_ai.shared.hash import ignore_keys, index_by, merge_dict
 
 
 class PermissionResource(DatabaseResource):
@@ -42,5 +42,42 @@ class PermissionResource(DatabaseResource):
             user_id=user.id if user else None,
         )), user, **kwargs)
 
+    @safe_db_query
+    def update(self, payload, **kwargs):
+        role_ids = [int(i) for i in payload.get('role_ids') or []]
+        roles_mapping = index_by(lambda x: x.id, self.roles or [])
+
+        role_ids_create = []
+        role_ids_delete = []
+
+        for role_id in role_ids:
+            if role_id not in roles_mapping:
+                role_ids_create.append(role_id)
+
+        for role_id in roles_mapping.keys():
+            if role_id not in role_ids:
+                role_ids_delete.append(role_id)
+
+        if role_ids_create:
+            db_connection.session.bulk_save_objects(
+                [RolePermission(
+                    permission_id=self.model.id,
+                    role_id=role_id,
+                    user_id=self.current_user.id,
+                ) for role_id in role_ids_create],
+                return_defaults=True,
+            )
+
+        if role_ids_delete:
+            delete_statement = RolePermission.__table__.delete().where(
+                RolePermission.permission_id == self.id,
+                RolePermission.role_id.in_(role_ids_delete),
+            )
+            db_connection.session.execute(delete_statement)
+
+        return super().update(ignore_keys(payload, [
+            'role_ids',
+        ]), **kwargs)
+
 
 PermissionResource.register_parent_model('role_id', Role)

mage_ai/api/resources/RoleResource.py

@@ -41,6 +41,7 @@ class RoleResource(DatabaseResource):
             roles = []
             for permission in permissions:
                 roles.append(permission.role)
+            roles = list(filter(lambda x: x, roles))
         else:
             roles = Role.query.all()