mage-ai 0.9.8__py3-none-any.whl → 0.9.9__py3-none-any.whl

mage_ai/data_preparation/models/block/__init__.py

@@ -515,7 +515,7 @@ class Block:
         if not replicated_block and \
                 (BlockType.DBT != block_type or BlockLanguage.YAML == language):
 
-            block_directory = f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
+            block_directory = self.file_directory_name(block_type)
             block_dir_path = os.path.join(repo_path, block_directory)
             if not os.path.exists(block_dir_path):
                 os.mkdir(block_dir_path)
@@ -567,6 +567,25 @@ class Block:
         )
         return block
 
+    @classmethod
+    def file_directory_name(self, block_type: BlockType) -> str:
+        return f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
+
+    @classmethod
+    def block_type_from_path(self, block_file_absolute_path: str) -> BlockType:
+        file_path = str(block_file_absolute_path).replace(get_repo_path(), '')
+        if file_path.startswith(os.sep):
+            file_path = file_path[1:]
+
+        file_path_parts = file_path.split(os.path.sep)
+        dir_name = file_path_parts[0]
+
+        for block_type in BlockType:
+            if BlockType.CUSTOM == block_type and dir_name == block_type:
+                return BlockType.CUSTOM
+            elif dir_name == f'{block_type}s':
+                return block_type
+
     @classmethod
     def get_all_blocks(self, repo_path) -> Dict:
         block_uuids = dict()

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

@@ -568,6 +568,7 @@ def config_file_loader_and_configuration(block, profile_target: str) -> Dict:
             SNOWFLAKE_DEFAULT_SCHEMA=profile.get('schema'),
             SNOWFLAKE_DEFAULT_WH=profile.get('warehouse'),
             SNOWFLAKE_USER=profile.get('user'),
+            SNOWFLAKE_ROLE=profile.get('role'),
         )
 
         if 'password' in profile:

mage_ai/data_preparation/models/custom_templates/custom_block_template.py

@@ -112,7 +112,7 @@ class CustomBlockTemplate(BaseConfig):
         language_to_use = language or self.language
         filename = '.'.join([
             self.template_uuid,
-            BLOCK_LANGUAGE_TO_FILE_EXTENSION[language_to_use],
+            BLOCK_LANGUAGE_TO_FILE_EXTENSION.get(language_to_use, ''),
         ])
 
         return File(
@@ -130,12 +130,17 @@ class CustomBlockTemplate(BaseConfig):
         if content:
             return Template(content).render(**(variables or {}))
 
-    def to_dict(self) -> Dict:
-        return merge_dict(self.to_dict_base(), dict(
+    def to_dict(self, include_content: bool = False) -> Dict:
+        data = merge_dict(self.to_dict_base(), dict(
             template_uuid=self.template_uuid,
             uuid=self.uuid,
         ))
 
+        if include_content:
+            data['content'] = self.load_template_content()
+
+        return data
+
     def to_dict_base(self) -> Dict:
         return dict(
             block_type=self.block_type,

mage_ai/data_preparation/models/custom_templates/utils.py

@@ -1,12 +1,13 @@
 import os
-from mage_ai.data_preparation.models.file import File
+from typing import Dict, List
+
 from mage_ai.data_preparation.models.custom_templates.constants import (
     CUSTOM_TEMPLATES_DIRECTORY,
     CUSTOM_TEMPLATES_DIRECTORY_ENVIRONMENT_VARIABLE,
 )
+from mage_ai.data_preparation.models.file import File
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.hash import group_by
-from typing import Dict, List
 
 
 def custom_templates_directory() -> str:
@@ -32,6 +33,8 @@ def flatten_files(
     children: List[Dict],
     parent_names: List[str] = None,
 ) -> List[Dict]:
+    if not children:
+        return []
     arr = []
 
     for child1 in children:
@@ -57,7 +60,7 @@ def group_and_hydrate_files(
 
     arr = []
 
-    for template_uuid, group in groups.items():
+    for template_uuid, _ in groups.items():
         custom_template = custom_template_class.load(template_uuid=template_uuid)
         if custom_template:
             arr.append(custom_template)

mage_ai/data_preparation/models/file.py

@@ -334,6 +334,9 @@ class File:
                 new_path,
             )
 
+        self.dir_path = dir_path
+        self.filename = filename
+
     def to_dict(self, include_content=False):
         data = dict(name=self.filename, path=os.path.join(self.dir_path, self.filename))
         if include_content:

mage_ai/data_preparation/models/pipeline.py

@@ -882,7 +882,11 @@ class Pipeline:
                     if block is None:
                         continue
                     if 'content' in block_data:
+                        from mage_ai.cache.block_action_object import BlockActionObjectCache
+
+                        cache_block_action_object = await BlockActionObjectCache.initialize_cache()
                         await block.update_content_async(block_data['content'], widget=widget)
+                        cache_block_action_object.update_block(block)
                     if 'callback_content' in block_data \
                             and block.callback_block:
                         await block.callback_block.update_content_async(
@@ -944,7 +948,12 @@ class Pipeline:
 
                         should_save_async = should_save_async or True
                     elif name and name != block.name:
+                        from mage_ai.cache.block_action_object import BlockActionObjectCache
+
+                        cache_block_action_object = await BlockActionObjectCache.initialize_cache()
+                        cache_block_action_object.update_block(block, remove=True)
                         block.update(extract(block_data, ['name']))
+                        cache_block_action_object.update_block(block)
                         block_uuid_mapping[block_data.get('uuid')] = block.uuid
                         should_save_async = should_save_async or True
 
@@ -1006,7 +1015,7 @@ class Pipeline:
         blocks_by_uuid,
         block,
         upstream_blocks,
-        priority=None,
+        priority: int = None,
     ):
         mapping = blocks_by_uuid.copy()
 

mage_ai/data_preparation/models/{project.py → project/__init__.py}

@@ -1,8 +1,10 @@
 import aiohttp
 
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.server.constants import VERSION
 from mage_ai.settings.repo import get_repo_path
+from typing import Dict
 
 
 class Project():
@@ -21,6 +23,21 @@ class Project():
     def project_uuid(self) -> str:
         return self.repo_config.project_uuid
 
+    @property
+    def openai_api_key(self) -> str:
+        return self.repo_config.openai_api_key
+
+    @property
+    def features(self) -> Dict:
+        data = {}
+        features = self.repo_config.features
+
+        for uuid in FeatureUUID:
+            key = uuid.value
+            data[key] = features.get(key) if features else None
+
+        return data
+
     async def latest_version(self) -> str:
         try:
             async with aiohttp.ClientSession() as session:

mage_ai/data_preparation/models/project/constants.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class FeatureUUID(str, Enum):
+    ADD_NEW_BLOCK_V2 = 'add_new_block_v2'

mage_ai/data_preparation/repo_manager.py

@@ -76,14 +76,15 @@ class RepoConfig:
                 repo_config.get('azure_container_instance_config')
             self.ecs_config = repo_config.get('ecs_config')
             self.emr_config = repo_config.get('emr_config')
+            self.features = repo_config.get('features', {})
             self.gcp_cloud_run_config = repo_config.get('gcp_cloud_run_config')
             self.k8s_executor_config = repo_config.get('k8s_executor_config')
             self.spark_config = repo_config.get('spark_config')
-
             self.notification_config = repo_config.get('notification_config', dict())
             self.queue_config = repo_config.get('queue_config', dict())
             self.project_uuid = repo_config.get('project_uuid')
             self.help_improve_mage = repo_config.get('help_improve_mage')
+            self.openai_api_key = repo_config.get('openai_api_key')
             self.retry_config = repo_config.get('retry_config')
 
             self.ldap_config = repo_config.get('ldap_config')
@@ -115,20 +116,22 @@ class RepoConfig:
 
     def to_dict(self, remote: bool = False) -> Dict:
         return dict(
-            project_type=self.project_type,
             azure_container_instance_config=self.azure_container_instance_config,
             ecs_config=self.ecs_config,
             emr_config=self.emr_config,
+            features=self.features,
             gcp_cloud_run_config=self.gcp_cloud_run_config,
+            help_improve_mage=self.help_improve_mage,
             notification_config=self.notification_config,
+            openai_api_key=self.openai_api_key,
+            project_type=self.project_type,
+            project_uuid=self.project_uuid,
             queue_config=self.queue_config,
+            remote_variables_dir=self.remote_variables_dir,
             repo_path=self.repo_path,
+            spark_config=self.spark_config,
             variables_dir=self.remote_variables_dir if remote else self.variables_dir,
             variables_retention_period=self.variables_retention_period,
-            remote_variables_dir=self.remote_variables_dir,
-            project_uuid=self.project_uuid,
-            help_improve_mage=self.help_improve_mage,
-            spark_config=self.spark_config,
         )
 
     def save(self, **kwargs) -> None: