mage-ai 0.9.69__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/api/resources/FileResource.py

@@ -148,7 +148,8 @@ class FileResource(GenericResource):
                 pipeline_file, pipeline_config = Pipeline.import_from_zip(content, overwrite)
                 tags = pipeline_config.get('tags', [])
                 pipeline_uuid = pipeline_config.get('uuid')
-                pipeline = Pipeline(pipeline_uuid, config=pipeline_config)
+                repo_path = get_repo_path(user=user)
+                pipeline = Pipeline(pipeline_uuid, config=pipeline_config, repo_path=repo_path)
                 if tags:
                     from mage_ai.cache.tag import TagCache
 
@@ -161,7 +162,7 @@ class FileResource(GenericResource):
 
                     block_cache = await BlockCache.initialize_cache()
                     for block in pipeline.blocks_by_uuid.values():
-                        block_cache.add_pipeline(block, pipeline)
+                        block_cache.add_pipeline(block, pipeline, repo_path)
 
                 return self(pipeline_file, user, **kwargs)
             else:
@@ -175,7 +176,9 @@ class FileResource(GenericResource):
                     overwrite=overwrite,
                 )
 
-                block_type = Block.block_type_from_path(dir_path)
+                block_type = Block.block_type_from_path(
+                    dir_path, repo_path=get_repo_path(user=user)
+                )
                 if block_type:
                     cache_block_action_object = await BlockActionObjectCache.initialize_cache()
                     cache_block_action_object.update_block(block_file_absolute_path=file.file_path)
@@ -235,7 +238,10 @@ class FileResource(GenericResource):
 
     @safe_db_query
     async def update(self, payload, **kwargs):
-        block_type = Block.block_type_from_path(self.model.dir_path)
+        block_type = Block.block_type_from_path(
+            self.model.dir_path,
+            repo_path=get_repo_path(user=self.current_user),
+        )
         cache_block_action_object = None
         if block_type:
             cache_block_action_object = await BlockActionObjectCache.initialize_cache()

mage_ai/api/resources/FileVersionResource.py

@@ -4,6 +4,7 @@ from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.file import File
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db import safe_db_query
+from mage_ai.settings.repo import get_repo_path
 
 
 class FileVersionResource(GenericResource):
@@ -21,7 +22,8 @@ class FileVersionResource(GenericResource):
             block_uuid = block_uuid[0]
 
         if pipeline_uuid and block_uuid:
-            pipeline = Pipeline.get(pipeline_uuid)
+            repo_path = get_repo_path(user=user)
+            pipeline = Pipeline.get(pipeline_uuid, repo_path=repo_path)
             if pipeline:
                 block = pipeline.get_block(block_uuid)
                 if block:

mage_ai/api/resources/GitBranchResource.py

@@ -1,16 +1,23 @@
 import os
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.data_preparation.git import REMOTE_NAME, Git, api
 from mage_ai.data_preparation.git.clients.base import Client as GitClient
-from mage_ai.data_preparation.git.utils import get_provider_from_remote_url
+from mage_ai.data_preparation.git.utils import (
+    get_oauth_access_token_for_user,
+    get_provider_from_remote_url,
+)
 from mage_ai.data_preparation.preferences import get_preferences
+from mage_ai.server.logger import Logger
+from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.path_fixer import remove_base_repo_path
 from mage_ai.shared.strings import capitalize_remove_underscore_lower
 
+logger = Logger().new_server_logger(__name__)
+
 
 def build_file_object(obj):
     arr = []
@@ -29,9 +36,21 @@ def build_file_object(obj):
 class GitBranchResource(GenericResource):
     @classmethod
     def get_git_manager(
-        self, user, setup_repo: bool = False, config_overwrite: Dict = None
+        self,
+        user=None,
+        context_data: Dict = None,
+        preferences=None,
+        repo_path: str = None,
+        setup_repo: bool = False,
+        config_overwrite: Dict = None,
     ) -> Git:
-        return Git.get_manager(setup_repo=setup_repo, user=user)
+        return Git.get_manager(
+            context_data=context_data,
+            preferences=preferences,
+            repo_path=repo_path,
+            setup_repo=setup_repo,
+            user=user,
+        )
 
     @classmethod
     def collection(self, query, meta, user, **kwargs):
@@ -53,7 +72,7 @@ class GitBranchResource(GenericResource):
             if remote_url:
                 provider = get_provider_from_remote_url(remote_url)
 
-            access_token = api.get_access_token_for_user(user, provider=provider)
+            access_token = get_oauth_access_token_for_user(user, provider=provider)
             if access_token:
                 branches = GitClient.get_client_for_provider(provider)(
                     access_token.token
@@ -64,7 +83,15 @@ class GitBranchResource(GenericResource):
             arr += [dict(name=branch) for branch in git_manager.branches]
 
             if include_remote_branches:
-                pass
+                try:
+                    git_manager.fetch()
+                    mage_remote = git_manager.origin
+                    arr += [
+                        dict(name=ref.name)
+                        for ref in mage_remote.refs
+                    ]
+                except Exception:
+                    logger.warning('Failed to fetch remote branches')
 
         return self.build_result_set(
             arr,
@@ -83,11 +110,22 @@ class GitBranchResource(GenericResource):
 
     @classmethod
     async def member(self, pk, user, **kwargs):
-        preferences = get_preferences(user=user)
+        context_data = kwargs.get('context_data')
+        repo_path = get_repo_path(context_data=context_data, user=user)
+        preferences = get_preferences(
+            repo_path=repo_path,
+            user=user,
+        )
         setup_repo = False
         if preferences.is_git_integration_enabled():
             setup_repo = True
-        git_manager = self.get_git_manager(user=user, setup_repo=setup_repo)
+        git_manager = self.get_git_manager(
+            user=user,
+            context_data=context_data,
+            preferences=preferences,
+            repo_path=repo_path,
+            setup_repo=setup_repo,
+        )
 
         display_format = kwargs.get('meta', {}).get('_format')
         if 'with_basic_details' == display_format:
@@ -98,7 +136,7 @@ class GitBranchResource(GenericResource):
                     modified_files=[],
                     name=git_manager.current_branch,
                     staged_files=[],
-                    sync_config=get_preferences().sync_config,
+                    sync_config=preferences.sync_config,
                     untracked_files=[],
                 ),
                 user,
@@ -116,39 +154,34 @@ class GitBranchResource(GenericResource):
                 modified_files=modified_files,
                 name=git_manager.current_branch,
                 staged_files=staged_files,
-                sync_config=get_preferences().sync_config,
+                sync_config=preferences.sync_config,
                 untracked_files=untracked_files,
             ),
             user,
             **kwargs,
         )
 
-    async def update(self, payload, **kwargs):
-        query = kwargs.get('query') or {}
-
-        git_manager = self.get_git_manager(user=self.current_user)
-        action_type = payload.get('action_type')
-        action_payload = payload.get('action_payload', dict())
-        action_remote = action_payload.get('remote', None)
-        action_branch = action_payload.get('branch', None)
-
-        files = payload.get('files', None)
-        message = payload.get('message', None)
-
+    @classmethod
+    def get_oauth_config(
+        cls,
+        remote_url: str = None,
+        remote_name: str = None,
+        user=None,
+    ) -> Tuple[str, str, str, Dict]:
+        git_manager = cls.get_git_manager(user=user)
         url = None
-        remote_url = query.get('remote_url', None)
         if remote_url:
-            url = remote_url[0]
-        elif action_remote:
-            remote = git_manager.repo.remotes[action_remote]
+            url = remote_url
+        elif remote_name:
+            remote = git_manager.repo.remotes[remote_name]
             url = list(remote.urls)[0]
 
         provider = ProviderName.GITHUB
         if url:
             provider = get_provider_from_remote_url(url)
 
-        access_token = api.get_access_token_for_user(
-            self.current_user, provider=provider
+        access_token = get_oauth_access_token_for_user(
+            user, provider=provider
         )
         http_access_token = git_manager.get_access_token()
 
@@ -158,9 +191,9 @@ class GitBranchResource(GenericResource):
             token = access_token.token
             user_from_api = api.get_user(token, provider=provider)
             # Default to mage user email if no email is returned from API
-            email = user_from_api.get(
-                'email', self.current_user.email if self.current_user else None
-            )
+            email = user_from_api.get('email')
+            if email is None and user:
+                email = user.email
             config_overwrite = dict(
                 username=user_from_api.get('username'),
                 email=email,
@@ -168,6 +201,37 @@ class GitBranchResource(GenericResource):
         elif http_access_token:
             token = http_access_token
 
+        return token, provider, url, config_overwrite
+
+    async def update(self, payload, **kwargs):
+        query = kwargs.get('query') or {}
+
+        git_manager = self.get_git_manager(user=self.current_user)
+        action_type = payload.get('action_type')
+        action_payload = payload.get('action_payload', dict())
+        action_remote = action_payload.get('remote', None)
+        action_branch = action_payload.get('branch', None)
+        action_arg = action_payload.get('arg', None)
+
+        # Eventually we would want to give the user the ability to pass in multiple
+        # arguments to the action, but for now we will just pass in a single argument.
+        action_kwargs = dict()
+        if action_arg:
+            action_kwargs[action_arg] = True
+
+        files = payload.get('files', None)
+        message = payload.get('message', None)
+
+        remote_url = query.get('remote_url', None)
+        if remote_url:
+            remote_url = remote_url[0]
+
+        token, provider, url, config_overwrite = self.get_oauth_config(
+            remote_url=remote_url,
+            remote_name=action_remote,
+            user=self.current_user,
+        )
+
         # Recreate git manager with updated config
         git_manager = self.get_git_manager(
             user=self.current_user, config_overwrite=config_overwrite
@@ -206,6 +270,7 @@ class GitBranchResource(GenericResource):
                             action_branch,
                             token,
                             config_overwrite=config_overwrite,
+                            **action_kwargs,
                         )
                         if custom_progress.other_lines:
                             lines = custom_progress.other_lines
@@ -236,6 +301,7 @@ class GitBranchResource(GenericResource):
                                 action_branch,
                                 token,
                                 config_overwrite=config_overwrite,
+                                **action_kwargs,
                             )
                         else:
                             custom_progress = api.fetch(
@@ -243,6 +309,7 @@ class GitBranchResource(GenericResource):
                                 url,
                                 token,
                                 config_overwrite=config_overwrite,
+                                **action_kwargs,
                             )
 
                         if custom_progress and custom_progress.other_lines:
@@ -271,6 +338,7 @@ class GitBranchResource(GenericResource):
                             url,
                             token,
                             config_overwrite=config_overwrite,
+                            **action_kwargs,
                         )
 
                         if custom_progress and custom_progress.other_lines:
@@ -304,6 +372,7 @@ class GitBranchResource(GenericResource):
                                 action_branch,
                                 token,
                                 config_overwrite=config_overwrite,
+                                **action_kwargs,
                             )
                         else:
                             self.model[
@@ -324,6 +393,7 @@ class GitBranchResource(GenericResource):
                             url,
                             token,
                             config_overwrite=config_overwrite,
+                            **action_kwargs,
                         )
                     else:
                         self.model[

mage_ai/api/resources/GitCustomBranchResource.py

@@ -1,7 +1,7 @@
 from typing import Dict, List
 
 from mage_ai.api.resources.GitBranchResource import GitBranchResource
-from mage_ai.data_preparation.git import REMOTE_NAME, Git
+from mage_ai.data_preparation.git import REMOTE_NAME, Git, api
 from mage_ai.data_preparation.sync import AuthType
 
 
@@ -17,6 +17,34 @@ class GitCustomBranchResource(GitBranchResource):
             user=user,
         )
 
+    @classmethod
+    def create(cls, payload, user, **kwargs):
+        branch = payload.get('name')
+        remote = payload.get('remote')
+
+        token, _, url, config_overwrite = cls.get_oauth_config(
+            remote_name=remote,
+            user=user,
+        )
+
+        git_manager = cls.get_git_manager(
+            user=user, config_overwrite=config_overwrite
+        )
+
+        if remote:
+            api.switch_branch(
+                remote,
+                url,
+                branch,
+                token,
+                config_overwrite=config_overwrite,
+                user=user,
+            )
+        else:
+            git_manager.switch_branch(branch)
+
+        return cls(dict(name=git_manager.current_branch), user, **kwargs)
+
     @classmethod
     async def member(self, pk, user, **kwargs):
         resource = await GitBranchResource.member(pk, user, **kwargs)

mage_ai/api/resources/GlobalDataProductResource.py

@@ -1,43 +1,80 @@
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.data_preparation.models.global_data_product import GlobalDataProduct
+from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.hash import ignore_keys
 
 
 class GlobalDataProductResource(GenericResource):
     @classmethod
     def collection(self, query, meta, user, **kwargs):
+        repo_path = query.get('repo_path', [None])
+        if repo_path:
+            repo_path = repo_path[0]
+
+        current_project = query.get('current_project', [False])
+        if current_project:
+            current_project = current_project[0]
+
+        if current_project:
+            repo_path = get_repo_path(user=user)
+
         return self.build_result_set(
-            sorted(GlobalDataProduct.load_all(), key=lambda x: x.uuid),
+            sorted(GlobalDataProduct.load_all(repo_path), key=lambda x: x.uuid),
             user,
             **kwargs,
         )
 
     @classmethod
     def create(self, payload, user, **kwargs):
+        repo_path = get_repo_path(user=user)
         uuid = payload.get('uuid')
-        if GlobalDataProduct.get(uuid):
+        if GlobalDataProduct.get(uuid, repo_path):
             error = ApiError.RESOURCE_INVALID.copy()
-            error.update(dict(message=f'A global data product with UUID {uuid} already exists.'))
+            error.update(
+                dict(message=f'A global data product with UUID {uuid} already exists.')
+            )
             raise ApiError(error)
 
-        model = GlobalDataProduct(uuid, **ignore_keys(payload, ['uuid']))
+        model = GlobalDataProduct(
+            uuid,
+            repo_path=repo_path,
+            **ignore_keys(payload, ['uuid']),
+        )
         model.save()
 
         return self(model, user, **kwargs)
 
     @classmethod
     def member(self, pk, user, **kwargs):
-        return self(GlobalDataProduct.get(pk), user, **kwargs)
+        query = kwargs.get('query', {})
+        project = query.get('project', [None])
+        if project:
+            project = project[0]
+        repo_path = get_repo_path(user=user)
+        return self(
+            GlobalDataProduct.get(
+                pk, repo_path=repo_path if not project else None, project=project
+            ),
+            user,
+            **kwargs,
+        )
 
     def delete(self, **kwargs):
         self.model.delete()
 
     def update(self, payload, **kwargs):
+        repo_path = get_repo_path(user=self.current_user)
         uuid = payload.get('uuid')
-        if self.model and self.model.uuid != uuid and GlobalDataProduct.get(uuid):
+        if (
+            self.model
+            and self.model.uuid != uuid
+            and GlobalDataProduct.get(uuid, repo_path)
+        ):
             error = ApiError.RESOURCE_INVALID.copy()
-            error.update(dict(message=f'A global data product with UUID {uuid} already exists.'))
+            error.update(
+                dict(message=f'A global data product with UUID {uuid} already exists.')
+            )
             raise ApiError(error)
 
         if self.model:

mage_ai/api/resources/GlobalHookResource.py

@@ -15,6 +15,7 @@ from mage_ai.data_preparation.models.global_hooks.models import (
 )
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.settings.repo import get_repo_path
 from mage_ai.settings.utils import base_repo_path
 from mage_ai.shared.array import find
 from mage_ai.shared.hash import ignore_keys
@@ -190,6 +191,8 @@ class GlobalHookResource(AsyncBaseResource):
 
 
 async def __load_pipelines(resource: GlobalHookResource):
+    repo_path = get_repo_path(user=resource.current_user)
+
     pipeline_uuids = []
     for res in resource.result_set():
         if res.model.pipeline_settings and res.model.pipeline_settings.get('uuid'):
@@ -199,7 +202,7 @@ async def __load_pipelines(resource: GlobalHookResource):
 
     async def get_pipeline(uuid):
         try:
-            return await Pipeline.get_async(uuid)
+            return await Pipeline.get_async(uuid, repo_path=repo_path)
         except Exception as err:
             err_message = f'Error loading pipeline {uuid}: {err}.'
             if err.__class__.__name__ == 'OSError' and 'Too many open files' in err.strerror:

mage_ai/api/resources/IntegrationDestinationResource.py

@@ -1,9 +1,12 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.data_integrations.destinations.constants import DESTINATIONS
 from mage_ai.data_preparation.models.constants import BlockType
-from mage_ai.data_preparation.models.pipelines.integration_pipeline import IntegrationPipeline
+from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
+    IntegrationPipeline,
+)
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.server.api.integration_sources import get_collection
+from mage_ai.settings.repo import get_repo_path
 
 
 class IntegrationDestinationResource(GenericResource):
@@ -27,7 +30,8 @@ class IntegrationDestinationResource(GenericResource):
         action_type = payload['action_type']
         if 'test_connection' == action_type:
             pipeline_uuid = payload['pipeline_uuid']
-            pipeline = IntegrationPipeline.get(pipeline_uuid)
+            repo_path = get_repo_path(user=user)
+            pipeline = IntegrationPipeline.get(pipeline_uuid, repo_path=repo_path)
             config = payload['config']
 
             try:

mage_ai/api/resources/IntegrationSourceResource.py

@@ -22,6 +22,7 @@ from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.schedules import PipelineRun
 from mage_ai.server.api.integration_sources import get_collection
+from mage_ai.settings.repo import get_repo_path
 
 
 def get_state_data_for_blocks(
@@ -140,13 +141,14 @@ class IntegrationSourceResource(GenericResource):
         success = False
         streams = []
 
+        repo_path = get_repo_path(user=user)
         action_type = payload['action_type']
         if 'test_connection' == action_type:
             pipeline_uuid = payload['pipeline_uuid']
             block_uuid = payload.get('block_uuid')
 
             if block_uuid:
-                pipeline = Pipeline.get(pipeline_uuid)
+                pipeline = Pipeline.get(pipeline_uuid, repo_path=repo_path)
                 block = pipeline.get_block(block_uuid)
                 try:
                     test_connection(block)
@@ -155,7 +157,7 @@ class IntegrationSourceResource(GenericResource):
                     traceback.print_exc()
                     error_message = str(e)
             else:
-                pipeline = IntegrationPipeline.get(pipeline_uuid)
+                pipeline = IntegrationPipeline.get(pipeline_uuid, repo_path=repo_path)
                 config = payload['config']
 
                 try:
@@ -166,7 +168,7 @@ class IntegrationSourceResource(GenericResource):
                     error_message = str(e)
         elif 'sample_data' == action_type:
             pipeline_uuid = payload['pipeline_uuid']
-            pipeline = IntegrationPipeline.get(pipeline_uuid)
+            pipeline = IntegrationPipeline.get(pipeline_uuid, repo_path=repo_path)
 
             streams_updated = pipeline.preview_data(
                 BlockType.DATA_LOADER,
@@ -222,4 +224,6 @@ class IntegrationSourceResource(GenericResource):
 
             return self(state_data, user, **kwargs)
 
-        return self(IntegrationPipeline.get(pk), user, **kwargs)
+        repo_path = get_repo_path(user=user)
+
+        return self(IntegrationPipeline.get(pk, repo_path=repo_path), user, **kwargs)

mage_ai/api/resources/IntegrationSourceStreamResource.py

@@ -1,10 +1,14 @@
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.models.pipelines.integration_pipeline import IntegrationPipeline
+from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
+    IntegrationPipeline,
+)
 from mage_ai.orchestration.db import safe_db_query
+from mage_ai.settings.repo import get_repo_path
 
 
 class IntegrationSourceStreamResource(GenericResource):
     @classmethod
     @safe_db_query
     def member(self, pk, user, **kwargs):
-        return self(IntegrationPipeline.get(pk), user, **kwargs)
+        repo_path = get_repo_path(user=user)
+        return self(IntegrationPipeline.get(pk, repo_path=repo_path), user, **kwargs)

mage_ai/api/resources/KernelProcessResource.py

@@ -0,0 +1,44 @@
+from mage_ai.api.errors import ApiError
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.kernels.models import KernelProcess
+from mage_ai.kernels.utils import get_process_info
+
+TERMINATE_INACTIVE = '__terminate_inactive__'
+
+
+class KernelProcessResource(GenericResource):
+    @classmethod
+    async def collection(cls, query, meta, user, **kwargs):
+        check_active_status = query.get('check_active_status', [False])
+        if check_active_status:
+            check_active_status = check_active_status[0]
+
+        return cls.build_result_set(
+            KernelProcess.load_all(check_active_status=check_active_status),
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    def member(self, pk, user, **kwargs):
+        if TERMINATE_INACTIVE == pk:
+            return self(KernelProcess.load(pid=TERMINATE_INACTIVE), user, **kwargs)
+
+        query = kwargs.get('query', {})
+        check_active_status = query.get('check_active_status', [False])
+        if check_active_status:
+            check_active_status = check_active_status[0]
+
+        info = get_process_info(pk, check_active_status=check_active_status)
+
+        if not info:
+            raise ApiError(ApiError.RESOURCE_NOT_FOUND)
+
+        kernel_process = KernelProcess.load(**info)
+        return self(kernel_process, user, **kwargs)
+
+    async def delete(self, **kwargs):
+        if self.model.pid == TERMINATE_INACTIVE:
+            KernelProcess.terminate_inactive()
+
+        self.model.terminate()