mage-ai 0.9.70__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

@@ -12,6 +12,7 @@ from mage_ai.data_preparation.git.utils import (
 )
 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
 
@@ -35,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):
@@ -97,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:
@@ -112,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,
@@ -130,7 +154,7 @@ 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,
@@ -167,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', user.email if 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,
@@ -187,6 +211,13 @@ class GitBranchResource(GenericResource):
         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)
@@ -239,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
@@ -269,6 +301,7 @@ class GitBranchResource(GenericResource):
                                 action_branch,
                                 token,
                                 config_overwrite=config_overwrite,
+                                **action_kwargs,
                             )
                         else:
                             custom_progress = api.fetch(
@@ -276,6 +309,7 @@ class GitBranchResource(GenericResource):
                                 url,
                                 token,
                                 config_overwrite=config_overwrite,
+                                **action_kwargs,
                             )
 
                         if custom_progress and custom_progress.other_lines:
@@ -304,6 +338,7 @@ class GitBranchResource(GenericResource):
                             url,
                             token,
                             config_overwrite=config_overwrite,
+                            **action_kwargs,
                         )
 
                         if custom_progress and custom_progress.other_lines:
@@ -337,6 +372,7 @@ class GitBranchResource(GenericResource):
                                 action_branch,
                                 token,
                                 config_overwrite=config_overwrite,
+                                **action_kwargs,
                             )
                         else:
                             self.model[
@@ -357,6 +393,7 @@ class GitBranchResource(GenericResource):
                             url,
                             token,
                             config_overwrite=config_overwrite,
+                            **action_kwargs,
                         )
                     else:
                         self.model[

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()

mage_ai/api/resources/KernelResource.py

@@ -1,4 +1,10 @@
+from datetime import datetime
+
 from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
+from mage_ai.kernels.models import KernelProcess, KernelWrapper
+from mage_ai.kernels.utils import find_ipykernel_launchers_info_async
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.server.active_kernel import (
     interrupt_kernel,
@@ -8,18 +14,31 @@ from mage_ai.server.active_kernel import (
 )
 from mage_ai.server.kernels import DEFAULT_KERNEL_NAME, KernelName, kernel_managers
 from mage_ai.services.ssh.aws.emr.utils import tunnel
+from mage_ai.shared.singletons.memory import get_memory_manager_controller
 
 
 class KernelResource(GenericResource):
     @classmethod
     @safe_db_query
-    def collection(self, query, meta, user, **kwargs):
+    async def collection(self, query, meta, user, **kwargs):
         kernels = []
 
+        if Project().is_feature_enabled(FeatureUUID.AUTOMATIC_KERNEL_CLEANUP):
+            # Only do this every minute
+            if int(datetime.utcnow().timestamp()) % 60 == 0:
+                kill_count, memory_freed = KernelProcess.terminate_inactive(
+                    await find_ipykernel_launchers_info_async(),
+                )
+                if kill_count >= 1:
+                    print(
+                        f'[KernelResource] Automatic kernel cleanup: {kill_count} '
+                        f'kernels terminated, {(memory_freed / 1024**3):.3f} GBs if memory freed.'
+                    )
+
         for kernel_name in KernelName:
             kernel = kernel_managers[kernel_name]
             if kernel.has_kernel:
-                kernels.append(kernel)
+                kernels.append(KernelWrapper(kernel))
 
         return self.build_result_set(
             kernels,
@@ -47,7 +66,7 @@ class KernelResource(GenericResource):
         if not kernel:
             kernel = kernel_managers[DEFAULT_KERNEL_NAME]
 
-        return self(kernel, user, **kwargs)
+        return self(KernelWrapper(kernel), user, **kwargs)
 
     @safe_db_query
     def update(self, payload, **kwargs):
@@ -56,13 +75,16 @@ class KernelResource(GenericResource):
         switch_active_kernel(self.model.kernel_name)
 
         if 'interrupt' == action_type:
+            get_memory_manager_controller().stop_all_events()
             interrupt_kernel()
         elif 'restart' == action_type:
             try:
+                get_memory_manager_controller().stop_all_events()
                 restart_kernel()
             except RuntimeError as e:
                 # RuntimeError: Cannot restart the kernel. No previous call to 'start_kernel'.
                 if 'start_kernel' in str(e):
+                    get_memory_manager_controller().stop_all_events()
                     start_kernel()
 
         def _callback(*args, **kwargs):

mage_ai/api/resources/OutputResource.py

@@ -9,20 +9,35 @@ from mage_ai.data_preparation.models.block.data_integration.utils import (
     read_data_from_cache,
 )
 from mage_ai.data_preparation.models.block.utils import serialize_output
+from mage_ai.data_preparation.models.constants import DATAFRAME_SAMPLE_COUNT
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.schedules import BlockRun
+from mage_ai.settings.repo import get_repo_path
 
 
 class OutputResource(GenericResource):
     @classmethod
     @safe_db_query
-    def collection(self, query, meta, user, **kwargs):
+    def collection(self, query_arg, meta, user, **kwargs):
+        query = {}
+        for key in [
+            'sample_count',
+        ]:
+            value = query_arg.get(key, [None])
+            if value:
+                value = value[0]
+                if value is not None:
+                    query[key] = value
+
         parent_model = kwargs['parent_model']
 
         outputs = []
         if type(parent_model) is BlockRun:
-            outputs = parent_model.get_outputs()
+            outputs = parent_model.get_outputs(
+                sample_count=query.get('sample_count', DATAFRAME_SAMPLE_COUNT),
+                exclude_blank_variable_uuids=True,
+            )
 
         return self.build_result_set(
             outputs,
@@ -45,7 +60,8 @@ class OutputResource(GenericResource):
         model = dict(outputs=[])
 
         if block_uuid and pipeline_uuid:
-            pipeline = Pipeline.get(pipeline_uuid)
+            repo_path = get_repo_path(user=user)
+            pipeline = Pipeline.get(pipeline_uuid, repo_path=repo_path)
             block = pipeline.get_block(block_uuid)
 
             if block.is_data_integration():
@@ -98,10 +114,12 @@ class OutputResource(GenericResource):
                             else:
                                 output_serialized = serialize_output(block, output)
 
-                        model['outputs'].append(dict(
-                            data=output_serialized,
-                            uuid=stream_id,
-                        ))
+                        model['outputs'].append(
+                            dict(
+                                data=output_serialized,
+                                uuid=stream_id,
+                            )
+                        )
 
                         if refresh and persist:
                             persist_data_for_stream(
@@ -134,6 +152,8 @@ class OutputResource(GenericResource):
         sample_count = payload.get('sample_count') or None
         if sample_count:
             sample_count = int(sample_count)
+        else:
+            sample_count = DATAFRAME_SAMPLE_COUNT
 
         model = dict(outputs=[])
 
@@ -163,9 +183,11 @@ class OutputResource(GenericResource):
                         if not output_serialized:
                             output_serialized = serialize_output(block, output)
 
-                        model['outputs'].append(dict(
-                            data=output_serialized,
-                            uuid=stream_id,
-                        ))
+                        model['outputs'].append(
+                            dict(
+                                data=output_serialized,
+                                uuid=stream_id,
+                            )
+                        )
 
         return self(model, user, **kwargs)