mage-ai 0.9.67__py3-none-any.whl → 0.9.68__py3-none-any.whl

mage_ai/api/monitors/BaseMonitor.py

@@ -24,8 +24,6 @@ class BaseMonitor:
         if self.error.type:
             data['type'] = self.error.type
 
-        loop = asyncio.get_event_loop()
-
         kwargs = dict(
             resource=self.resource,
             **extract(
@@ -48,6 +46,7 @@ class BaseMonitor:
             ),
         )
 
+        loop = asyncio.get_event_loop()
         if loop is not None:
             loop.create_task(
                 UsageStatisticLogger().error(

mage_ai/api/policies/PipelinePolicy.py

@@ -85,9 +85,11 @@ PipelinePolicy.allow_write([
     'conditionals',
     'clone_pipeline_uuid',
     'custom_template_uuid',
+    'description',
     'extensions',
     'llm',
     'name',
+    'tags',
     'type',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/resources/BlockLayoutItemResource.py

@@ -6,7 +6,7 @@ import urllib.parse
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.widget.constants import ChartType
 from mage_ai.data_preparation.variable_manager import get_global_variables
@@ -73,7 +73,7 @@ class BlockLayoutItemResource(GenericResource):
                 data_source_type = data_source_config.get('type')
                 pipeline_uuid = data_source_config.get('pipeline_uuid')
 
-                block = Block.get_block(
+                block = BlockFactory.get_block(
                     block_config.get('name') or file_path or uuid,
                     block_uuid,
                     block_type,

mage_ai/api/resources/BlockResource.py

@@ -12,6 +12,7 @@ from mage_ai.cache.block_action_object.constants import (
     OBJECT_TYPE_MAGE_TEMPLATE,
 )
 from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.block.utils import clean_name
 from mage_ai.data_preparation.models.constants import (
     FILE_EXTENSION_TO_BLOCK_LANGUAGE,
@@ -163,7 +164,7 @@ class BlockResource(GenericResource):
             if payload.get('type') == BlockType.DBT and language == BlockLanguage.SQL and content:
                 from mage_ai.data_preparation.models.block.dbt import DBTBlock
 
-                dbt_block = DBTBlock(
+                dbt_block = DBTBlock.create(
                     name,
                     clean_name(name),
                     BlockType.DBT,
@@ -323,7 +324,7 @@ class BlockResource(GenericResource):
         if BlockType.DBT == block_type:
             from mage_ai.data_preparation.models.block.dbt import DBTBlock
 
-            block = DBTBlock(
+            block = DBTBlock.create(
                 block_uuid,
                 block_uuid,
                 block_type,
@@ -331,7 +332,7 @@ class BlockResource(GenericResource):
                 language=language,
             )
         else:
-            block = Block.get_block(block_uuid, block_uuid, block_type, language=language)
+            block = BlockFactory.get_block(block_uuid, block_uuid, block_type, language=language)
 
         if not block.exists():
             error.update(ApiError.RESOURCE_NOT_FOUND)

mage_ai/api/resources/CacheItemResource.py

@@ -50,7 +50,7 @@ class CacheItemResource(AsyncBaseResource):
         )
 
         if CacheItemType.DBT == item_type:
-            block = DBTBlock(
+            block = DBTBlock.create(
                 pk,
                 pk,
                 BlockType.DBT,

mage_ai/api/resources/GitBranchResource.py

@@ -152,14 +152,10 @@ class GitBranchResource(GenericResource):
         )
         http_access_token = git_manager.get_access_token()
 
+        config_overwrite = None
         token = None
         if access_token:
             token = access_token.token
-        elif http_access_token:
-            token = http_access_token
-
-        config_overwrite = None
-        if 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(
@@ -169,6 +165,8 @@ class GitBranchResource(GenericResource):
                 username=user_from_api.get('username'),
                 email=email,
             )
+        elif http_access_token:
+            token = http_access_token
 
         # Recreate git manager with updated config
         git_manager = self.get_git_manager(

mage_ai/api/resources/PageBlockLayoutResource.py

@@ -3,7 +3,7 @@ import urllib.parse
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.presenters.block_layout.page import PageBlockLayout
 from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 from mage_ai.shared.utils import clean_name
@@ -52,7 +52,7 @@ class PageBlockLayoutResource(GenericResource):
                     if len(parts) >= 2:
                         block_uuid_use = os.path.join(*parts[1:])
 
-                block = Block.get_block(
+                block = BlockFactory.get_block(
                     block_uuid_use,
                     block_uuid_use,
                     block_type,

mage_ai/api/resources/PipelineResource.py

@@ -362,6 +362,8 @@ class PipelineResource(BaseResource):
         clone_pipeline_uuid = payload.get('clone_pipeline_uuid')
         template_uuid = payload.get('custom_template_uuid')
         name = payload.get('name')
+        description = payload.get('description')
+        tags = payload.get('tags')
         pipeline_type = payload.get('type')
         llm_payload = payload.get('llm')
         pipeline = None
@@ -375,8 +377,10 @@ class PipelineResource(BaseResource):
         else:
             pipeline = Pipeline.create(
                 name,
+                description=description,
                 pipeline_type=pipeline_type,
                 repo_path=get_repo_path(),
+                tags=tags,
             )
 
             if llm_payload:
@@ -451,6 +455,15 @@ class PipelineResource(BaseResource):
             cache = await PipelineCache.initialize_cache()
             cache.add_model(resource.model)
 
+            tags = resource.model.tags
+            if tags:
+                from mage_ai.cache.tag import TagCache
+
+                cache = await TagCache.initialize_cache()
+
+                for tag_uuid in tags:
+                    cache.add_pipeline(tag_uuid, resource.model)
+
         self.on_create_callback = _on_create_callback
 
         return self(pipeline, user, **kwargs)

mage_ai/api/resources/PipelineRunResource.py

@@ -12,7 +12,7 @@ from mage_ai.data_preparation.models.triggers import (
     ScheduleStatus,
     ScheduleType,
 )
-from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.models.schedules import (
     BlockRun,
     PipelineRun,
@@ -241,6 +241,15 @@ class PipelineRunResource(DatabaseResource):
         has_next = results_size > limit
         final_end_idx = results_size - 1 if has_next else results_size
 
+        # Expire pipeline runs that are in progress so that the latest status is returned
+        # the next time they are fetched.
+        for run in results:
+            if run.status in (
+                PipelineRun.PipelineRunStatus.RUNNING,
+                PipelineRun.PipelineRunStatus.INITIAL,
+            ):
+                db_connection.session.expire(run)
+
         result_set = self.build_result_set(
             results[0:final_end_idx],
             user,

mage_ai/cache/tag.py

@@ -144,6 +144,9 @@ class TagCache(BaseCache):
 
         mapping[key][KEY_FOR_PIPELINES] = pipelines_dict
 
+        if not pipelines_dict and mapping[key].keys() == set([KEY_FOR_PIPELINES]):
+            mapping.pop(key, None)
+
         self.set(self.cache_key, mapping)
 
     async def initialize_cache_for_all_objects(self) -> None:

mage_ai/cluster_manager/manage.py

@@ -70,11 +70,14 @@ def get_workspaces(cluster_type: ClusterType, **kwargs) -> List[Workspace]:
     repo_path = get_repo_path()
     projects_folder = os.path.join(repo_path, 'projects')
     if is_main_project:
+        # sort by file last modified time
         projects = [
-            f.name.split('.')[0]
+            f
             for f in os.scandir(projects_folder)
             if not f.is_dir() and f.name.endswith('.yaml')
         ]
+        projects.sort(key=lambda f: os.path.getmtime(f.path), reverse=True)
+        projects = [os.path.splitext(f.name)[0] for f in projects]
     else:
         instances = get_instances(cluster_type, **kwargs)
         projects = [instance.get('name') for instance in instances]

mage_ai/data_preparation/executors/k8s_block_executor.py

@@ -1,6 +1,9 @@
 from typing import Dict
 
+from jinja2 import Template
+
 from mage_ai.data_preparation.executors.block_executor import BlockExecutor
+from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.constants import DEFAULT_NAMESPACE
 from mage_ai.services.k8s.job_manager import JobManager as K8sJobManager
@@ -29,9 +32,13 @@ class K8sBlockExecutor(BlockExecutor):
             job_name_prefix = self.executor_config.job_name_prefix
 
         if self.executor_config.namespace:
-            namespace = self.executor_config.namespace
+            namespace = Template(self.executor_config.namespace).render(
+                variables=lambda x: global_vars.get(x) if global_vars else None,
+                **get_template_vars()
+            )
         else:
             namespace = DEFAULT_NAMESPACE
+
         job_manager = K8sJobManager(
             job_name=f'mage-{job_name_prefix}-block-{block_run_id}',
             logger=self.logger,

mage_ai/data_preparation/executors/k8s_pipeline_executor.py

@@ -1,8 +1,11 @@
 import traceback
 from typing import Dict
 
+from jinja2 import Template
+
 from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.constants import DEFAULT_NAMESPACE
 from mage_ai.services.k8s.job_manager import JobManager as K8sJobManager
@@ -39,6 +42,7 @@ class K8sPipelineExecutor(PipelineExecutor):
     ) -> None:
         job_manager = self.get_job_manager(
             pipeline_run_id=pipeline_run_id,
+            global_vars=global_vars,
             **kwargs,
         )
         cmd = self._run_commands(
@@ -54,15 +58,22 @@ class K8sPipelineExecutor(PipelineExecutor):
     def get_job_manager(
         self,
         pipeline_run_id: int = None,
+        global_vars: Dict = None,
         **kwargs,
     ) -> K8sJobManager:
+        if global_vars is None:
+            global_vars = dict()
         if not self.executor_config.job_name_prefix:
             job_name_prefix = 'data-prep'
         else:
             job_name_prefix = self.executor_config.job_name_prefix
 
         if self.executor_config.namespace:
-            namespace = self.executor_config.namespace
+
+            namespace = Template(self.executor_config.namespace).render(
+                variables=lambda x: global_vars.get(x) if global_vars else None,
+                **get_template_vars()
+            )
         else:
             namespace = DEFAULT_NAMESPACE
 

mage_ai/data_preparation/logging/gcs_logger_manager.py

@@ -10,9 +10,9 @@ from mage_ai.shared.config import BaseConfig
 
 @dataclass
 class GCSConfig(BaseConfig):
-    path_to_credentials: str
     bucket: str
     prefix: str
+    path_to_credentials: str = None
 
 
 class GCSLoggerManager(LoggerManager):
@@ -23,9 +23,12 @@ class GCSLoggerManager(LoggerManager):
     ):
         super().__init__(repo_config=repo_config, **kwargs)
         self.gcs_config = GCSConfig.load(config=self.logging_config.destination_config)
-        credentials = service_account.Credentials.from_service_account_file(
-            self.gcs_config.path_to_credentials
-        )
+        if self.gcs_config.path_to_credentials:
+            credentials = service_account.Credentials.from_service_account_file(
+                self.gcs_config.path_to_credentials
+            )
+        else:
+            credentials = None
         self.gcs_client = storage.Client(credentials=credentials)
 
     def create_log_filepath_dir(self, path):

mage_ai/data_preparation/models/block/__init__.py

@@ -90,7 +90,7 @@ from mage_ai.data_preparation.templates.data_integrations.utils import get_templ
 from mage_ai.data_preparation.templates.template import load_template
 from mage_ai.server.kernel_output_parser import DataType
 from mage_ai.services.spark.config import SparkConfig
-from mage_ai.services.spark.spark import get_spark_session
+from mage_ai.services.spark.spark import SPARK_ENABLED, get_spark_session
 from mage_ai.settings.platform.constants import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.array import unique_by
@@ -723,41 +723,6 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                 widget=widget,
             )
 
-    @classmethod
-    def block_class_from_type(self, block_type: str, language=None, pipeline=None) -> 'Block':
-        from mage_ai.data_preparation.models.block.constants import BLOCK_TYPE_TO_CLASS
-        from mage_ai.data_preparation.models.block.integration import (
-            DestinationBlock,
-            SourceBlock,
-            TransformerBlock,
-        )
-        from mage_ai.data_preparation.models.block.r import RBlock
-        from mage_ai.data_preparation.models.block.sql import SQLBlock
-        from mage_ai.data_preparation.models.widget import Widget
-
-        if BlockType.CHART == block_type:
-            return Widget
-        elif BlockType.DBT == block_type:
-            from mage_ai.data_preparation.models.block.dbt import DBTBlock
-
-            return DBTBlock
-        elif pipeline and PipelineType.INTEGRATION == pipeline.type:
-            if BlockType.CALLBACK == block_type:
-                return CallbackBlock
-            elif BlockType.CONDITIONAL == block_type:
-                return ConditionalBlock
-            elif BlockType.DATA_LOADER == block_type:
-                return SourceBlock
-            elif BlockType.DATA_EXPORTER == block_type:
-                return DestinationBlock
-            else:
-                return TransformerBlock
-        elif BlockLanguage.SQL == language:
-            return SQLBlock
-        elif BlockLanguage.R == language:
-            return RBlock
-        return BLOCK_TYPE_TO_CLASS.get(block_type)
-
     @classmethod
     def create(
         self,
@@ -777,6 +742,8 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
         widget: bool = False,
         downstream_block_uuids: List[str] = None,
     ) -> 'Block':
+        from mage_ai.data_preparation.models.block.block_factory import BlockFactory
+
         """
         1. Create a new folder for block_type if not exist
         2. Create a new python file with code template
@@ -870,7 +837,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                     language,
                 )
 
-        block = self.block_class_from_type(block_type, pipeline=pipeline)(
+        block = BlockFactory.block_class_from_type(
+            block_type,
+            language=language,
+            pipeline=pipeline,
+        )(
             name,
             uuid,
             block_type,
@@ -931,34 +902,6 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                     block_uuids[t.value].append(f.split('.')[0])
         return block_uuids
 
-    @classmethod
-    def get_block(
-        self,
-        name,
-        uuid,
-        block_type,
-        configuration=None,
-        content=None,
-        language=None,
-        pipeline=None,
-        status=BlockStatus.NOT_EXECUTED,
-    ) -> 'Block':
-        block_class = self.block_class_from_type(
-            block_type,
-            language=language,
-            pipeline=pipeline,
-        ) or Block
-        return block_class(
-            name,
-            uuid,
-            block_type,
-            configuration=configuration,
-            content=content,
-            language=language,
-            pipeline=pipeline,
-            status=status,
-        )
-
     @classmethod
     def get_block_from_file_path(self, file_path: str) -> 'Block':
         parts = get_path_parts(file_path)
@@ -2445,8 +2388,10 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
             return Template(self.executor_type).render(**get_template_vars())
         return self.executor_type
 
-    def get_pipelines_from_cache(self) -> List[Dict]:
-        arr = BlockCache().get_pipelines(self)
+    def get_pipelines_from_cache(self, block_cache: BlockCache = None) -> List[Dict]:
+        if block_cache is None:
+            block_cache = BlockCache()
+        arr = block_cache.get_pipelines(self)
 
         return unique_by(
             arr,
@@ -2497,11 +2442,13 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
     def to_dict(
         self,
         include_block_catalog: bool = False,
+        include_block_pipelines: bool = False,
         include_callback_blocks: bool = False,
         include_content: bool = False,
         include_outputs: bool = False,
         include_outputs_spark: bool = False,
         sample_count: int = None,
+        block_cache: BlockCache = None,
         check_if_file_exists: bool = False,
         **kwargs,
     ) -> Dict:
@@ -2515,6 +2462,9 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
         if include_block_catalog and self.is_data_integration() and self.pipeline:
             data['catalog'] = self.get_catalog_from_file()
 
+        if include_block_pipelines:
+            data['pipelines'] = self.get_pipelines_from_cache(block_cache=block_cache)
+
         if include_outputs:
             include_outputs_use = include_outputs
             if self.is_using_spark() and self.compute_management_enabled():
@@ -2551,6 +2501,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
         include_outputs: bool = False,
         include_outputs_spark: bool = False,
         sample_count: int = None,
+        block_cache: BlockCache = None,
         check_if_file_exists: bool = False,
         **kwargs,
     ) -> Dict:
@@ -2596,7 +2547,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
             data['tags'] = self.tags()
 
         if include_block_pipelines:
-            data['pipelines'] = self.get_pipelines_from_cache()
+            data['pipelines'] = self.get_pipelines_from_cache(block_cache=block_cache)
 
         return data
 
@@ -3093,6 +3044,8 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
         return global_vars
 
     def get_spark_session(self):
+        if not SPARK_ENABLED:
+            return None
         if self.spark_init and (not self.pipeline or
                                 not self.pipeline.spark_config):
             return self.spark

mage_ai/data_preparation/models/block/block_factory.py

@@ -0,0 +1,77 @@
+from mage_ai.data_preparation.models.block import Block, CallbackBlock, ConditionalBlock
+from mage_ai.data_preparation.models.block.constants import BLOCK_TYPE_TO_CLASS
+from mage_ai.data_preparation.models.block.integration import (
+    DestinationBlock,
+    SourceBlock,
+    TransformerBlock,
+)
+from mage_ai.data_preparation.models.block.r import RBlock
+from mage_ai.data_preparation.models.block.sql import SQLBlock
+from mage_ai.data_preparation.models.constants import (
+    BlockLanguage,
+    BlockStatus,
+    BlockType,
+    PipelineType,
+)
+from mage_ai.data_preparation.models.widget import Widget
+
+try:
+    from mage_ai.data_preparation.models.block.dbt.block_sql import DBTBlockSQL
+    from mage_ai.data_preparation.models.block.dbt.block_yaml import DBTBlockYAML
+except Exception:
+    print('DBT library not installed.')
+
+
+class BlockFactory:
+    @classmethod
+    def block_class_from_type(self, block_type: str, language=None, pipeline=None) -> 'Block':
+        if BlockType.CHART == block_type:
+            return Widget
+        elif BlockType.DBT == block_type:
+            if language == BlockLanguage.YAML:
+                return DBTBlockYAML
+            return DBTBlockSQL
+        elif pipeline and PipelineType.INTEGRATION == pipeline.type:
+            if BlockType.CALLBACK == block_type:
+                return CallbackBlock
+            elif BlockType.CONDITIONAL == block_type:
+                return ConditionalBlock
+            elif BlockType.DATA_LOADER == block_type:
+                return SourceBlock
+            elif BlockType.DATA_EXPORTER == block_type:
+                return DestinationBlock
+            else:
+                return TransformerBlock
+        elif BlockLanguage.SQL == language:
+            return SQLBlock
+        elif BlockLanguage.R == language:
+            return RBlock
+        return BLOCK_TYPE_TO_CLASS.get(block_type)
+
+    @classmethod
+    def get_block(
+        self,
+        name,
+        uuid,
+        block_type,
+        configuration=None,
+        content=None,
+        language=None,
+        pipeline=None,
+        status=BlockStatus.NOT_EXECUTED,
+    ) -> 'Block':
+        block_class = BlockFactory.block_class_from_type(
+            block_type,
+            language=language,
+            pipeline=pipeline,
+        ) or Block
+        return block_class(
+            name,
+            uuid,
+            block_type,
+            configuration=configuration,
+            content=content,
+            language=language,
+            pipeline=pipeline,
+            status=status,
+        )

mage_ai/data_preparation/models/block/dbt/block.py

@@ -26,19 +26,17 @@ from mage_ai.shared.parsers import encode_complex
 
 
 class DBTBlock(Block):
-    def __new__(cls, *args, **kwargs) -> 'DBTBlock':
+    @classmethod
+    def create(cls, *args, **kwargs) -> 'DBTBlock':
         """
         Factory for the child blocks
         """
         # Import Child blocks here to prevent cycle import
         from mage_ai.data_preparation.models.block.dbt.block_sql import DBTBlockSQL
         from mage_ai.data_preparation.models.block.dbt.block_yaml import DBTBlockYAML
-        if cls is DBTBlock:
-            if kwargs.get('language', BlockLanguage.SQL) == BlockLanguage.YAML:
-                return super(DBTBlock, cls).__new__(DBTBlockYAML)
-            return super(DBTBlock, cls).__new__(DBTBlockSQL)
-        else:
-            return super(DBTBlock, cls).__new__(cls, *args, **kwargs)
+        if kwargs.get('language', BlockLanguage.SQL) == BlockLanguage.YAML:
+            return DBTBlockYAML(*args, **kwargs)
+        return DBTBlockSQL(*args, **kwargs)
 
     @property
     def base_project_path(self) -> Union[str, os.PathLike]:

mage_ai/data_preparation/models/block/global_data_product/__init__.py

@@ -2,12 +2,15 @@ from logging import Logger
 from typing import Dict, List
 
 from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.global_data_product import GlobalDataProduct
-from mage_ai.orchestration.triggers import global_data_product as trigger
 
 
 class GlobalDataProductBlock(Block):
-    def get_global_data_product(self) -> GlobalDataProduct:
+    def get_global_data_product(self):
+        # Avoid circular dependency of Pipeline class
+        from mage_ai.data_preparation.models.global_data_product import (
+            GlobalDataProduct,
+        )
+
         override_configuration = (self.configuration or {}).get('global_data_product', {})
         global_data_product = GlobalDataProduct.get(override_configuration.get('uuid'))
 
@@ -31,6 +34,9 @@ class GlobalDataProductBlock(Block):
         logging_tags: Dict = None,
         **kwargs,
     ) -> List:
+        # Avoid circular dependency of Pipeline class
+        from mage_ai.orchestration.triggers import global_data_product as trigger
+
         trigger.trigger_and_check_status(
             self.get_global_data_product(),
             block=self,