mage-ai 0.8.97__py3-none-any.whl → 0.8.99__py3-none-any.whl

mage_ai/api/policies/AutocompleteItemPolicy.py

@@ -2,12 +2,13 @@ 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.presenters.AutocompleteItemPresenter import AutocompleteItemPresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class AutocompleteItemPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 AutocompleteItemPolicy.allow_actions([

mage_ai/api/policies/BasePolicy.py

@@ -15,7 +15,7 @@ from mage_ai.api.utils import (
     has_at_least_viewer_role,
     is_owner,
 )
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 from mage_ai.services.tracking.metrics import increment
 from mage_ai.settings import DISABLE_NOTEBOOK_EDIT_ACCESS, REQUIRE_USER_AUTHENTICATION
@@ -43,7 +43,7 @@ class BasePolicy():
 
     @property
     def entity(self) -> Tuple[Union[Permission.Entity, None], Union[str, None]]:
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
     @classmethod
     def action_rule(self, action):

mage_ai/api/policies/BlockTemplatePolicy.py

@@ -2,12 +2,13 @@ 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.presenters.BlockTemplatePresenter import BlockTemplatePresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class BlockTemplatePolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 BlockTemplatePolicy.allow_actions([

mage_ai/api/policies/ClusterPolicy.py

@@ -2,12 +2,13 @@ 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.presenters.ClusterPresenter import ClusterPresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class ClusterPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 ClusterPolicy.allow_actions([

mage_ai/api/policies/DataProviderPolicy.py

@@ -2,12 +2,13 @@ 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.presenters.DataProviderPresenter import DataProviderPresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class DataProviderPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 DataProviderPolicy.allow_actions([

mage_ai/api/policies/EventRulePolicy.py

@@ -2,12 +2,13 @@ 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.presenters.EventRulePresenter import EventRulePresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class EventRulePolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 EventRulePolicy.allow_actions([

mage_ai/api/policies/ExtensionOptionPolicy.py

@@ -2,12 +2,13 @@ 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.presenters.ExtensionOptionPresenter import ExtensionOptionPresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class ExtensionOptionPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 ExtensionOptionPolicy.allow_actions([

mage_ai/api/policies/FileVersionPolicy.py

@@ -2,12 +2,13 @@ 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.presenters.FileVersionPresenter import FileVersionPresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class FileVersionPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 FileVersionPolicy.allow_actions([

mage_ai/api/policies/GitBranchPolicy.py

@@ -53,3 +53,12 @@ GitBranchPolicy.allow_write(GitBranchPresenter.default_attributes + [
 ], on_action=[
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
+
+GitBranchPolicy.allow_query([
+    'include_remote_branches',
+    'repository',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/KernelPolicy.py

@@ -2,12 +2,13 @@ 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.presenters.KernelPresenter import KernelPresenter
+from mage_ai.orchestration.db.models.oauth import Permission
 
 
 class KernelPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 KernelPolicy.allow_actions([

mage_ai/api/policies/LogPolicy.py

@@ -3,7 +3,7 @@ from mage_ai.api.operations import constants
 from mage_ai.api.policies.BasePolicy import BasePolicy
 from mage_ai.api.presenters.LogPresenter import LogPresenter
 from mage_ai.data_preparation.models.pipeline import Pipeline
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 from mage_ai.orchestration.db.models.schedules import BlockRun
 
@@ -18,7 +18,7 @@ class LogPolicy(BasePolicy):
             elif issubclass(parent_model.__class__, Pipeline):
                 return Permission.Entity.PIPELINE, parent_model.uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 LogPolicy.allow_actions([

mage_ai/api/policies/OauthPolicy.py

@@ -10,9 +10,16 @@ class OauthPolicy(BasePolicy):
 
 OauthPolicy.allow_actions([
     constants.CREATE,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+OauthPolicy.allow_actions([
     constants.DETAIL,
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
+    OauthScope.CLIENT_PUBLIC,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
 
@@ -20,6 +27,13 @@ OauthPolicy.allow_read(OauthPresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.CREATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+OauthPolicy.allow_read(OauthPresenter.default_attributes, scopes=[
+    OauthScope.CLIENT_PRIVATE,
+    OauthScope.CLIENT_PUBLIC,
+], on_action=[
     constants.DETAIL,
 ], condition=lambda policy: policy.has_at_least_viewer_role())
 
@@ -40,4 +54,5 @@ OauthPolicy.allow_query([
     constants.DETAIL,
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
+    OauthScope.CLIENT_PUBLIC,
 ], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/OutputPolicy.py

@@ -2,7 +2,7 @@ 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.presenters.OutputPresenter import OutputPresenter
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 from mage_ai.orchestration.db.models.schedules import BlockRun
 
@@ -15,7 +15,7 @@ class OutputPolicy(BasePolicy):
             if type(parent_model) is BlockRun:
                 return Permission.Entity.PIPELINE, parent_model.pipeline_run.pipeline_uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 OutputPolicy.allow_actions([

mage_ai/api/policies/PipelinePolicy.py

@@ -2,7 +2,7 @@ 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.presenters.PipelinePresenter import PipelinePresenter
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 
 
@@ -12,7 +12,7 @@ class PipelinePolicy(BasePolicy):
         if self.resource and self.resource.model:
             return Permission.Entity.PIPELINE, self.resource.model.uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 PipelinePolicy.allow_actions([

mage_ai/api/policies/PipelineRunPolicy.py

@@ -2,7 +2,7 @@ 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.presenters.PipelineRunPresenter import PipelineRunPresenter
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 
 
@@ -19,7 +19,7 @@ class PipelineRunPolicy(BasePolicy):
         if self.resource and self.resource.model:
             return Permission.Entity.PIPELINE, self.resource.model.pipeline_uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 PipelineRunPolicy.allow_actions([

mage_ai/api/policies/PipelineSchedulePolicy.py

@@ -2,7 +2,7 @@ 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.presenters.PipelineSchedulePresenter import PipelineSchedulePresenter
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 
 
@@ -12,7 +12,7 @@ class PipelineSchedulePolicy(BasePolicy):
         if self.resource and self.resource.model:
             return Permission.Entity.PIPELINE, self.resource.model.pipeline_uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 PipelineSchedulePolicy.allow_actions([

mage_ai/api/policies/PullRequestPolicy.py

@@ -0,0 +1,64 @@
+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.presenters.PullRequestPresenter import PullRequestPresenter
+
+
+class PullRequestPolicy(BasePolicy):
+    pass
+
+
+PullRequestPolicy.allow_actions([
+    constants.DETAIL,
+    constants.LIST,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+PullRequestPolicy.allow_actions([
+    constants.DETAIL,
+    constants.LIST,
+    constants.CREATE,
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], condition=lambda policy: policy.has_at_least_editor_role())
+
+
+PullRequestPolicy.allow_read(PullRequestPresenter.default_attributes + [
+], scopes=[
+    OauthScope.CLIENT_PRIVATE
+], on_action=[
+    constants.DETAIL,
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+
+PullRequestPolicy.allow_read(PullRequestPresenter.default_attributes + [
+], scopes=[
+    OauthScope.CLIENT_PRIVATE
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_editor_role())
+
+
+PullRequestPolicy.allow_write([
+    'base_branch',
+    'body',
+    'compare_branch',
+    'repository',
+    'title',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_editor_role())
+
+
+PullRequestPolicy.allow_query([
+    'repository',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.LIST,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/SessionPolicy.py

@@ -1,12 +1,13 @@
 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.orchestration.db.models.oauth import Permission
 
 
 class SessionPolicy(BasePolicy):
     @property
     def entity(self):
-        return None, None
+        return Permission.Entity.ANY, None
 
 
 SessionPolicy.allow_actions([
@@ -45,6 +46,8 @@ SessionPolicy.allow_write([
     'email',
     'password',
     'username',
+    'token',
+    'provider',
 ], scopes=[
     OauthScope.CLIENT_PUBLIC,
 ], on_action=[

mage_ai/api/policies/VariablePolicy.py

@@ -2,7 +2,7 @@ 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.presenters.VariablePresenter import VariablePresenter
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 
 
@@ -13,7 +13,7 @@ class VariablePolicy(BasePolicy):
         if parent_model:
             return Permission.Entity.PIPELINE, parent_model.uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 VariablePolicy.allow_actions([

mage_ai/api/policies/WidgetPolicy.py

@@ -2,7 +2,7 @@ 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.presenters.WidgetPresenter import WidgetPresenter
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Permission
 
 
@@ -13,7 +13,7 @@ class WidgetPolicy(BasePolicy):
         if parent_model:
             return Permission.Entity.PIPELINE, parent_model.uuid
 
-        return Permission.Entity.PROJECT, get_repo_path()
+        return Permission.Entity.PROJECT, get_project_uuid()
 
 
 WidgetPolicy.allow_actions([

mage_ai/api/policies/WorkspacePolicy.py

@@ -33,9 +33,7 @@ WorkspacePolicy.allow_read(WorkspacePresenter.default_attributes, scopes=[
     constants.LIST,
 ], condition=lambda policy: policy.has_at_least_admin_role())
 
-WorkspacePolicy.allow_read(WorkspacePresenter.default_attributes + [
-    'success',
-], scopes=[
+WorkspacePolicy.allow_read(WorkspacePresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.CREATE,
@@ -50,6 +48,8 @@ WorkspacePolicy.allow_write([
     'storage_class_name',
     'service_account_name',
     'container_config',
+    'storage_access_mode',
+    'storage_request_size',
     'cluster_name',
     'task_definition',
     'container_name',

mage_ai/api/presenters/PipelinePresenter.py

@@ -12,6 +12,7 @@ class PipelinePresenter(BasePresenter):
         'executor_count',
         'executor_type',
         'name',
+        'notification_config',
         'spark_config',
         'type',
         'updated_at',

mage_ai/api/presenters/PullRequestPresenter.py

@@ -0,0 +1,16 @@
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class PullRequestPresenter(BasePresenter):
+    default_attributes = [
+        'body',
+        'created_at',
+        'id',
+        'is_merged',
+        'last_modified',
+        'merged',
+        'state',
+        'title',
+        'url',
+        'user',
+    ]

mage_ai/api/presenters/StatusPresenter.py

@@ -10,4 +10,6 @@ class StatusPresenter(BasePresenter):
         'instance_type',
         'disable_pipeline_edit_access',
         'require_user_authentication',
+        'project_type',
+        'project_uuid',
     ]

mage_ai/api/presenters/SyncPresenter.py

@@ -12,6 +12,7 @@ class SyncPresenter(BasePresenter):
         'ssh_private_key',
         'ssh_public_key',
         'sync_on_pipeline_run',
+        'sync_on_start',
         'type',
         'user_git_settings',
         'username',

mage_ai/api/presenters/WorkspacePresenter.py

@@ -8,4 +8,6 @@ class WorkspacePresenter(BasePresenter):
         'cluster_type',
         'access',
         'repo_path',
+        'project_uuid',
+        'success',
     ]

mage_ai/api/resources/GitBranchResource.py

@@ -1,9 +1,11 @@
+import os
+from github import Auth, Github
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.data_preparation.git import Git
+from mage_ai.data_preparation.git import api
 from mage_ai.data_preparation.preferences import get_preferences
 from typing import Dict, List
-import os
 
 
 def build_file_object(obj):
@@ -21,9 +23,34 @@ def build_file_object(obj):
 class GitBranchResource(GenericResource):
     @classmethod
     def collection(self, query, meta, user, **kwargs):
-        git_manager = Git.get_manager(user=user)
+        arr = []
+
+        include_remote_branches = query.get('include_remote_branches', None)
+        if include_remote_branches:
+            include_remote_branches = include_remote_branches[0]
+
+        repository = query.get('repository', None)
+        if repository:
+            repository = repository[0]
+
+            access_token = api.get_access_token_for_user(user)
+            if access_token:
+                auth = Auth.Token(access_token.token)
+                g = Github(auth=auth)
+                repo = g.get_repo(repository)
+                branches = repo.get_branches()
+
+                for branch in branches:
+                    arr.append(dict(name=branch.name))
+        else:
+            git_manager = Git.get_manager(user=user)
+            arr += [dict(name=branch) for branch in git_manager.branches]
+
+            if include_remote_branches:
+                pass
+
         return self.build_result_set(
-            [dict(name=branch) for branch in git_manager.branches],
+            arr,
             user,
             **kwargs,
         )
@@ -112,38 +139,66 @@ class GitBranchResource(GenericResource):
             push = payload.get('push', None)
 
             if push and 'remote' in push and 'branch' in push:
-                import git
+                from git.exc import GitCommandError
 
                 try:
-                    custom_progress = git_manager.push_remote_branch(
-                        push['remote'],
-                        push['branch'],
-                    )
-                    if custom_progress.other_lines:
-                        lines = custom_progress.other_lines
-                        if type(lines) is list:
-                            lines = '\n'.join(lines)
-                        self.model['progress'] = lines
-                except git.exc.GitCommandError as err:
+                    access_token = api.get_access_token_for_user(self.current_user)
+                    if access_token:
+                        remote = git_manager.repo.remotes[push['remote']]
+                        custom_progress = api.push(
+                            remote.name,
+                            [url for url in remote.urls][0],
+                            push['branch'],
+                            access_token.token,
+                        )
+                        if custom_progress.other_lines:
+                            lines = custom_progress.other_lines
+                            if type(lines) is list:
+                                lines = '\n'.join(lines)
+                            self.model['progress'] = lines
+                    else:
+                        self.model['error'] = \
+                            'Please authenticate with GitHub before trying to push.'
+                except GitCommandError as err:
                     self.model['error'] = str(err)
             else:
                 git_manager.push()
         elif action_type == 'pull':
             pull = payload.get('pull', None)
             if pull and 'remote' in pull:
-                import git
+                from git.exc import GitCommandError
 
                 try:
-                    custom_progress = git_manager.pull_remote_branch(
-                        pull['remote'],
-                        pull.get('branch'),
-                    )
-                    if custom_progress.other_lines:
-                        lines = custom_progress.other_lines
-                        if type(lines) is list:
-                            lines = '\n'.join(lines)
-                        self.model['progress'] = lines
-                except git.exc.GitCommandError as err:
+                    access_token = api.get_access_token_for_user(self.current_user)
+                    if access_token:
+                        branch_name = payload.get('branch')
+                        remote = git_manager.repo.remotes[pull['remote']]
+                        url = list(remote.urls)[0]
+
+                        custom_progress = None
+                        if branch_name:
+                            custom_progress = api.pull(
+                                remote.name,
+                                url,
+                                pull.get('branch'),
+                                access_token.token,
+                            )
+                        else:
+                            custom_progress = api.fetch(
+                                remote.name,
+                                url,
+                                access_token.token,
+                            )
+
+                        if custom_progress and custom_progress.other_lines:
+                            lines = custom_progress.other_lines
+                            if type(lines) is list:
+                                lines = '\n'.join(lines)
+                            self.model['progress'] = lines
+                    else:
+                        self.model['error'] = \
+                            'Please authenticate with GitHub before trying to pull.'
+                except GitCommandError as err:
                     self.model['error'] = str(err)
             else:
                 git_manager.pull()
@@ -223,4 +278,4 @@ class GitBranchResource(GenericResource):
 
     def remotes(self, limit: int = None) -> List[Dict]:
         git_manager = Git.get_manager(user=self.current_user)
-        return git_manager.remotes(limit=limit)
+        return git_manager.remotes(limit=limit, user=self.current_user)