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

mage_ai/api/constants.py

@@ -1,10 +1,11 @@
 from enum import Enum
 
-META_KEY_LIMIT = '_limit'
-META_KEY_OFFSET = '_offset'
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class AttributeOperationType(str, Enum):
+    QUERY = 'query'
     READ = 'read'
     WRITE = 'write'
 
@@ -17,3 +18,37 @@ class AuthorizeStatusType(str, Enum):
     ALL = 'all'
     FAILED = 'failed'
     SUCCEEDED = 'succeeded'
+
+
+OPERATION_TYPE_TO_ACCESS_MAPPING = {
+  OperationType.ALL: Permission.Access.OPERATION_ALL,
+  OperationType.CREATE: Permission.Access.CREATE,
+  OperationType.DELETE: Permission.Access.DELETE,
+  OperationType.DETAIL: Permission.Access.DETAIL,
+  OperationType.LIST: Permission.Access.LIST,
+  OperationType.UPDATE: Permission.Access.UPDATE,
+}
+
+OPERATION_TYPE_DISABLE_TO_ACCESS_MAPPING = {
+  OperationType.ALL: Permission.Access.DISABLE_OPERATION_ALL,
+  OperationType.CREATE: Permission.Access.DISABLE_CREATE,
+  OperationType.DELETE: Permission.Access.DISABLE_DELETE,
+  OperationType.DETAIL: Permission.Access.DISABLE_DETAIL,
+  OperationType.LIST: Permission.Access.DISABLE_LIST,
+  OperationType.UPDATE: Permission.Access.DISABLE_UPDATE,
+}
+
+ATTRIBUTE_OPERATION_TYPE_TO_ACCESS_MAPPING = {
+  AttributeOperationType.QUERY: Permission.Access.QUERY,
+  AttributeOperationType.READ: Permission.Access.READ,
+  AttributeOperationType.WRITE: Permission.Access.WRITE,
+}
+
+ATTRIBUTE_OPERATION_TYPE_DISABLE_TO_ACCESS_MAPPING = {
+  AttributeOperationType.QUERY: Permission.Access.DISABLE_QUERY,
+  AttributeOperationType.READ: Permission.Access.DISABLE_READ,
+  AttributeOperationType.WRITE: Permission.Access.DISABLE_WRITE,
+}
+
+META_KEY_LIMIT = '_limit'
+META_KEY_OFFSET = '_offset'

mage_ai/api/errors.py

@@ -1,3 +1,6 @@
+from typing import Dict
+
+
 class ApiError(Exception):
     UNAUTHORIZED_ACCESS = {
         'code': 403,
@@ -37,8 +40,9 @@ class ApiError(Exception):
         'type': 'invalid_oauth_token',
     }
 
-    def __init__(self, opts={}):
-        self.code = opts.get('code')
-        self.errors = opts.get('errors')
-        self.message = opts.get('message')
-        self.type = opts.get('type')
+    def __init__(self, opts: Dict = None):
+        if opts:
+            self.code = opts.get('code')
+            self.errors = opts.get('errors')
+            self.message = opts.get('message')
+            self.type = opts.get('type')

mage_ai/api/middleware.py

@@ -18,8 +18,8 @@ from mage_ai.shared.array import find
 
 
 class OAuthMiddleware(RequestHandler):
-    def initialize(self, bypass_oauth_check: bool = False) -> None:
-        self.bypass_oauth_check = bypass_oauth_check
+    def initialize(self, **kwargs) -> None:
+        self.bypass_oauth_check = kwargs.get('bypass_oauth_check', False)
 
     def prepare(self):
         self.request.__setattr__('current_user', None)

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:
@@ -255,11 +260,19 @@ class BaseOperation():
             return await self.__resource_class().process_create(
                 payload,
                 self.user,
+                result_set_from_external=policy.result_set(),
                 **options,
             )
         elif LIST == self.action:
-            def _build_authorize_query(parsed_value: Any, policy=policy) -> Callable:
-                return policy.authorize_query(parsed_value)
+            def _build_authorize_query(
+                parsed_value: Any,
+                policy=policy,
+                updated_options=updated_options,
+                **_kwargs,
+            ) -> Callable:
+                return policy.authorize_query(parsed_value, **ignore_keys(updated_options, [
+                    'query',
+                ]))
 
             if parser:
                 value_parsed, parser_found, error = await parser.parse_query_and_authorize(
@@ -279,6 +292,7 @@ class BaseOperation():
                 self.query,
                 self.meta,
                 self.user,
+                result_set_from_external=policy.result_set(),
                 **options,
             )
 
@@ -307,8 +321,15 @@ class BaseOperation():
         if DELETE == self.action:
             await res.process_delete(**updated_options)
         elif DETAIL == self.action:
-            def _build_authorize_query(parsed_value: Any, policy=policy) -> Callable:
-                return policy.authorize_query(parsed_value)
+            def _build_authorize_query(
+                parsed_value: Any,
+                policy=policy,
+                updated_options=updated_options,
+                **_kwargs,
+            ) -> Callable:
+                return policy.authorize_query(parsed_value, **ignore_keys(updated_options, [
+                    'query',
+                ]))
 
             if parser:
                 value_parsed, parser_found, error = await parser.parse_query_and_authorize(
@@ -327,7 +348,9 @@ class BaseOperation():
                             **merge_dict(updated_options, dict(query=self.query)),
                         )
                         policy = self.__policy_class()(res, self.user, **updated_options)
-                        await policy.authorize_query(self.query)
+                        await policy.authorize_query(self.query, **ignore_keys(updated_options, [
+                            'query',
+                        ]))
                     else:
                         raise error
             else:
@@ -351,8 +374,15 @@ class BaseOperation():
             else:
                 await _build_authorize_attributes(payload)
 
-            def _build_authorize_query(parsed_value: Any, policy=policy) -> Callable:
-                return policy.authorize_query(parsed_value)
+            def _build_authorize_query(
+                parsed_value: Any,
+                policy=policy,
+                updated_options=updated_options,
+                **_kwargs,
+            ) -> Callable:
+                return policy.authorize_query(parsed_value, **ignore_keys(updated_options, [
+                    'query',
+                ]))
 
             if parser:
                 value_parsed, parser_found, error = await parser.parse_query_and_authorize(
@@ -362,7 +392,7 @@ class BaseOperation():
                 )
                 self.query = value_parsed
             else:
-                await _build_authorize_query(self.query)
+                await _build_authorize_query(self.query, **updated_options)
 
             options = merge_dict(updated_options.copy(), dict(query=self.query))
             options.pop('payload', None)

mage_ai/api/operations/constants.py

@@ -1,5 +1,6 @@
 from enum import Enum
 
+ALL = 'all'
 CREATE = 'create'
 DELETE = 'delete'
 DETAIL = 'detail'
@@ -19,9 +20,9 @@ COOKIE_PREFIX = '__COOKIE__'
 
 
 class OperationType(str, Enum):
-    ALL = 'all'
-    CREATE = 'create'
-    DELETE = 'delete'
-    DETAIL = 'detail'
-    LIST = 'list'
-    UPDATE = 'update'
+    ALL = ALL
+    CREATE = CREATE
+    DELETE = DELETE
+    DETAIL = DETAIL
+    LIST = LIST
+    UPDATE = UPDATE

mage_ai/api/parsers/BaseParser.py

@@ -53,7 +53,7 @@ class BaseParser:
     @classmethod
     def parse_query(
         self,
-        parser: Callable[['BaseParser', Any, Dict], Any],
+        parser: Callable[['BaseParser', Any, Union[ApiError], Dict], Any],
         on_action: Union[List[OperationType], OperationType] = None,
         on_authorize_status: Union[List[AuthorizeStatusType], AuthorizeStatusType] = None,
         scopes: List[OauthScopeType] = None,
@@ -83,7 +83,7 @@ class BaseParser:
     @classmethod
     def parse_read(
         self,
-        parser: Callable[['BaseParser', Any, Dict], Any],
+        parser: Callable[['BaseParser', Any, Union[ApiError, Dict]], Any],
         on_action: Union[List[OperationType], OperationType] = None,
         on_authorize_status: Union[List[AuthorizeStatusType], AuthorizeStatusType] = None,
         scopes: List[OauthScopeType] = None,
@@ -113,7 +113,7 @@ class BaseParser:
     @classmethod
     def parse_write(
         self,
-        parser: Callable[['BaseParser', Any, Dict], Any],
+        parser: Callable[['BaseParser', Any, Union[ApiError, Dict]], Any],
         on_action: Union[List[OperationType], OperationType] = None,
         on_authorize_status: Union[List[AuthorizeStatusType], AuthorizeStatusType] = None,
         scopes: List[OauthScopeType] = None,
@@ -161,6 +161,7 @@ class BaseParser:
             self.query_parsers.get(self.__class__.__name__) or {},
             value,
             authorize_status,
+            error=error,
             **kwargs,
         )
 
@@ -242,6 +243,7 @@ class BaseParser:
         parser_mapping: Dict,
         value: Any,
         authorize_status: AuthorizeStatusType,
+        error: ApiError = None,
         **kwargs,
     ) -> Tuple[Any, bool]:
         api_operation_action = self.options.get('api_operation_action') or \
@@ -262,7 +264,7 @@ class BaseParser:
         if not parser:
             return value, False
 
-        value_parsed = parser(self, value, **kwargs)
+        value_parsed = parser(self, value, authorize_status=authorize_status, error=error, **kwargs)
 
         if value_parsed and inspect.isawaitable(value_parsed):
             value_parsed = await value_parsed

mage_ai/api/parsers/PipelineScheduleParser.py

@@ -8,6 +8,8 @@ from mage_ai.api.policies.PipelineSchedulePolicy import (
     READABLE_ATTRIBUTES_FOR_PIPELINE_INTERACTIONS,
     WRITABLE_ATTRIBUTES_FOR_PIPELINE_INTERACTIONS,
 )
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.shared.hash import extract
 
 
@@ -23,8 +25,39 @@ async def parse_read_update(parser, _value: Dict, **kwargs) -> Dict:
     return {}
 
 
-async def parse_write(parser, value: Dict, **kwargs) -> Dict:
-    return extract(value, WRITABLE_ATTRIBUTES_FOR_PIPELINE_INTERACTIONS)
+async def parse_write(
+    parser,
+    value: Dict,
+    authorize_status: AuthorizeStatusType = None,
+    **kwargs,
+) -> Dict:
+    if AuthorizeStatusType.FAILED == authorize_status:
+        parsed_value = extract(value or {}, WRITABLE_ATTRIBUTES_FOR_PIPELINE_INTERACTIONS)
+    else:
+        parsed_value = (value or {}).copy()
+
+    pipeline = parser.policy.parent_model()
+    if not pipeline:
+        pipeline = parser.policy.resource.pipeline
+
+    if Project(pipeline.repo_config if pipeline else None).is_feature_enabled(
+        FeatureUUID.INTERACTIONS,
+    ):
+        pipeline_interaction = (parser.policy.result_set().context.data or {}).get(
+            'pipeline_interactions',
+            {},
+        ).get(pipeline.uuid)
+
+        if pipeline_interaction:
+            variables = parsed_value.get('variables')
+            if pipeline_interaction and variables:
+                variables_allowed = await pipeline_interaction.variables()
+                parsed_value['variables'] = extract(
+                    variables or {},
+                    (variables_allowed or {}).keys(),
+                )
+
+    return parsed_value
 
 
 async def parse_write_update(parser, value: Dict, **kwargs) -> Dict:
@@ -66,6 +99,7 @@ PipelineScheduleParser.parse_write(
     ],
     on_authorize_status=[
         AuthorizeStatusType.FAILED,
+        AuthorizeStatusType.SUCCEEDED,
     ],
     scopes=[
         OauthScopeType.CLIENT_PRIVATE,

mage_ai/api/policies/BackfillPolicy.py

@@ -81,6 +81,7 @@ BackfillPolicy.allow_query([
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
 BackfillPolicy.allow_query([
+    'include_run_count',
     'pipeline_uuid',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,