mage-ai 0.9.55__py3-none-any.whl → 0.9.57__py3-none-any.whl

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -731,9 +731,7 @@ class WorkloadManager:
 
         return env_vars
 
-    def __get_configurable_parameters(
-        self, workspace_config: KubernetesWorkspaceConfig
-    ) -> Dict:
+    def get_default_values(self) -> Dict:
         service_account_name_default = None
         storage_class_name_default = None
         storage_access_mode_default = None
@@ -757,19 +755,31 @@ class WorkloadManager:
         except Exception:
             pass
 
+        return dict(
+            service_account_name=service_account_name_default,
+            storage_class_name=storage_class_name_default,
+            storage_access_mode=storage_access_mode_default,
+            storage_request_size=storage_request_size_default,
+        )
+
+    def __get_configurable_parameters(
+        self, workspace_config: KubernetesWorkspaceConfig
+    ) -> Dict:
+        default_values = self.get_default_values()
+
         storage_request_size = workspace_config.storage_request_size
         if storage_request_size is None:
-            storage_request_size = storage_request_size_default
+            storage_request_size = default_values.get('storage_request_size')
         else:
             storage_request_size = f'{storage_request_size}Gi'
 
         return dict(
             service_account_name=workspace_config.service_account_name
-            or service_account_name_default,
+            or default_values.get('service_account_name'),
             storage_class_name=workspace_config.storage_class_name
-            or storage_class_name_default,
+            or default_values.get('storage_class_name'),
             storage_access_mode=workspace_config.storage_access_mode
-            or storage_access_mode_default,
+            or default_values.get('storage_access_mode'),
             storage_request_size=storage_request_size,
         )
 

mage_ai/data_preparation/models/block/__init__.py

@@ -12,9 +12,11 @@ from contextlib import contextmanager, redirect_stderr, redirect_stdout
 from datetime import datetime
 from inspect import Parameter, isfunction, signature
 from logging import Logger
+from pathlib import Path
 from queue import Queue
 from typing import Any, Callable, Dict, Generator, List, Set, Tuple, Union
 
+import inflection
 import pandas as pd
 import simplejson
 import yaml
@@ -57,6 +59,7 @@ from mage_ai.data_preparation.models.constants import (
     DATAFRAME_ANALYSIS_MAX_COLUMNS,
     DATAFRAME_ANALYSIS_MAX_ROWS,
     DATAFRAME_SAMPLE_COUNT_PREVIEW,
+    FILE_EXTENSION_TO_BLOCK_LANGUAGE,
     NON_PIPELINE_EXECUTABLE_BLOCK_TYPES,
     BlockColor,
     BlockLanguage,
@@ -84,6 +87,10 @@ from mage_ai.shared.environments import get_env, is_debug
 from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 from mage_ai.shared.logger import BlockFunctionExec
 from mage_ai.shared.parsers import encode_complex
+from mage_ai.shared.path_fixer import (
+    add_root_repo_path_to_relative_path,
+    get_path_parts,
+)
 from mage_ai.shared.strings import format_enum
 from mage_ai.shared.utils import clean_name as clean_name_orig
 from mage_ai.shared.utils import is_spark_env
@@ -586,7 +593,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
     def file_path(self) -> str:
         file_path = self.get_file_path_from_source()
         if file_path:
-            return file_path
+            return add_root_repo_path_to_relative_path(file_path)
 
         return self.__build_file_path(
             self.repo_path or os.getcwd(),
@@ -753,16 +760,28 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
         )
         file_is_from_another_project = (
             file_path_from_source and
-            from_another_project(file_path_from_source)
+            from_another_project(
+                file_path=file_path_from_source,
+                other_file_path=pipeline.dir_path if pipeline else None,
+            )
         )
+        absolute_file_path = add_root_repo_path_to_relative_path(
+            file_path_from_source,
+        ) if file_path_from_source else None
+
+        if not file_is_from_another_project and \
+                (not absolute_file_path or not os.path.exists(absolute_file_path)):
 
-        if not file_is_from_another_project:
             if not replicated_block and \
                     (BlockType.DBT != block_type or BlockLanguage.YAML == language) and \
                     BlockType.GLOBAL_DATA_PRODUCT != block_type:
 
                 block_directory = self.file_directory_name(block_type)
-                block_dir_path = os.path.join(repo_path, block_directory)
+                if absolute_file_path:
+                    block_dir_path = os.path.dirname(absolute_file_path)
+                else:
+                    block_dir_path = os.path.join(repo_path, block_directory)
+
                 if not os.path.exists(block_dir_path):
                     os.mkdir(block_dir_path)
                     with open(os.path.join(block_dir_path, '__init__.py'), 'w'):
@@ -895,6 +914,41 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
             status=status,
         )
 
+    @classmethod
+    def get_block_from_file_path(self, file_path: str) -> 'Block':
+        parts = get_path_parts(file_path)
+
+        if parts and len(parts) >= 3:
+            # If file_path == transformers/test4.py
+            # parts ==
+            # ('/home/src/default_repo/default_platform2/project3', 'transformers', 'test4.py')
+
+            # If project platform platform activated, then parts ==
+            # ('/home/src', 'default_repo', 'data_loaders/astral_violet.py')
+
+            root_project_full_path, path, file_path_base = parts
+
+            if project_platform_activated():
+                # ('data_loaders', 'astral_violet.py')
+                file_parts = Path(file_path_base).parts
+                block_type = inflection.singularize(str(file_parts[0]))
+                block_uuid = str(Path(*file_parts[1:]).with_suffix(''))
+            else:
+                block_type = inflection.singularize(str(path))
+                block_uuid = str(Path(file_path_base).with_suffix(''))
+
+            extension = Path(file_path).suffix.replace('.', '')
+            configuration = dict(file_path=file_path, file_source=dict(path=file_path))
+            language = FILE_EXTENSION_TO_BLOCK_LANGUAGE.get(extension)
+
+            return self.get_block(
+                block_uuid,
+                block_uuid,
+                block_type,
+                configuration=configuration,
+                language=language,
+            )
+
     def all_upstream_blocks_completed(
         self,
         completed_block_uuids: Set[str],
@@ -945,6 +999,8 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                 if len(pipelines) == 0:
                     os.remove(self.file_path)
             return
+
+        # TODO (tommy dang): delete this block from all pipelines in all projects
         # If pipeline is not specified, delete the block from all pipelines and delete the file.
         pipelines = Pipeline.get_pipelines_by_block(self, widget=widget)
         if not force:
@@ -954,6 +1010,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                         f'Block {self.uuid} has downstream dependencies in pipeline {p.uuid}. '
                         'Please remove the dependencies before deleting the block.'
                     )
+
         for p in pipelines:
             p.delete_block(
                 p.get_block(self.uuid, widget=widget),
@@ -961,6 +1018,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                 commit=commit,
                 force=force,
             )
+
         os.remove(self.file_path)
 
     def execute_with_callback(
@@ -2336,7 +2394,15 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
             self.__update_pipeline_block(widget=widget)
         return self
 
-    async def update_content_async(self, content, widget=False) -> 'Block':
+    async def update_content_async(
+        self,
+        content,
+        error_if_file_missing: bool = True,
+        widget: bool = False,
+    ) -> 'Block':
+        if error_if_file_missing and not self.file.exists():
+            raise Exception(f'File for block {self.uuid} does not exist at {self.file.file_path}.')
+
         block_content = await self.content_async()
         if content != block_content:
             self.status = BlockStatus.UPDATED

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

@@ -65,7 +65,9 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
             if file_path:
                 return file_path
 
-        file_path = self.configuration.get('file_path')
+        file_path = self.configuration.get('file_path') or ''
+        if file_path and Path(file_path).parts[0] == 'dbt':
+            file_path = os.path.join(*Path(file_path).parts[1:]) or ''
 
         return str((Path(self.repo_path)) / DBT_DIRECTORY_NAME / file_path)
 
@@ -82,9 +84,15 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
             if project_path:
                 return project_path
 
+        try:
+            diff = Path(self.file_path).relative_to(self.base_project_path)
+            if diff:
+                diff = Path(diff).parts[0]
+        except ValueError:
+            diff = self.file_path
+
         return str(
-            Path(self.base_project_path) /
-            self.configuration.get('file_path', '').split(os.sep)[0]
+            Path(self.base_project_path) / diff,
         )
 
     @property
@@ -187,6 +195,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
         # Get upstream nodes via dbt list
         with Profiles(self.project_path, self.pipeline.variables) as profiles:
             try:
+
                 args = [
                     'list',
                     # project-dir
@@ -218,6 +227,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
                             pipeline=self.pipeline,
                             uuid=self.uuid,
                         ),
+                        hydrate_configuration=False,
                     )
                 ]
 

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

@@ -199,7 +199,21 @@ class DBTAdapter:
         profiles_path = self.__profiles.interpolate()
 
         # set dbt flags
-        user_config = read_user_config(profiles_path)
+        # Need to add profiles.yml file
+        try:
+            user_config = read_user_config(profiles_path)
+        except Exception as err:
+            print(f'[ERROR] DBTAdapter.open: {err}.')
+
+            if not profiles_path.endswith('profiles.yaml') and \
+                    not profiles_path.endswith('profiles.yml'):
+
+                try:
+                    user_config = read_user_config(os.path.join(profiles_path, 'profiles.yml'))
+                except Exception as err2:
+                    print(f'[ERROR] DBTAdapter.open: {err2}.')
+                    raise err
+
         adapter_config = DBTAdapterConfig(
             project_dir=self.project_path,
             profiles_dir=profiles_path,

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

@@ -135,19 +135,21 @@ class Profiles(object):
 
         # write interpolated profiles.yml
         interpoalted_profiles_full_path = interpolated_profiles_dir / PROFILES_FILE_NAME
-        if os.path.exists(interpoalted_profiles_full_path):
-            try:
-                with interpoalted_profiles_full_path.open('w') as f:
-                    yaml.safe_dump(self.profiles, f)
-            except Exception as e:
-                raise ProfilesError(
-                    f'Failed to write interpolated `{PROFILES_FILE_NAME}`' +
-                    f' to `{interpolated_profiles_dir}`: {e}'
-                )
+        if not os.path.exists(interpoalted_profiles_full_path):
+            os.makedirs(os.path.dirname(interpoalted_profiles_full_path), exist_ok=True)
+
+        try:
+            with interpoalted_profiles_full_path.open('w') as f:
+                yaml.safe_dump(self.profiles, f)
+        except Exception as e:
+            raise ProfilesError(
+                f'Failed to write interpolated `{PROFILES_FILE_NAME}`' +
+                f' to `{interpolated_profiles_dir}`: {e}'
+            )
 
-            self.__interpolated_profiles_dir = str(interpolated_profiles_dir)
-            self.is_interpolated = True
-            return self.__interpolated_profiles_dir
+        self.__interpolated_profiles_dir = str(interpolated_profiles_dir)
+        self.is_interpolated = True
+        return self.__interpolated_profiles_dir
 
     def __del__(self) -> None:
         self.clean()

mage_ai/data_preparation/models/block/platform/mixins.py

@@ -79,7 +79,10 @@ class ProjectPlatformAccessible:
         return config
 
     def is_from_another_project(self) -> bool:
-        return self.project_platform_activated and from_another_project(self.__file_source_path())
+        return self.project_platform_activated and from_another_project(
+            self.__file_source_path(),
+            other_file_path=self.pipeline.dir_path if self.pipeline else None,
+        )
 
     def get_file_path_from_source(self) -> str:
         if not self.project_platform_activated:
@@ -201,6 +204,7 @@ class ProjectPlatformAccessible:
         block_class,
         block_dict,
         node: Dict = None,
+        hydrate_configuration: bool = True,
     ):
         block_type = block_dict['block_type']
         configuration = block_dict['configuration'] or {}
@@ -209,7 +213,7 @@ class ProjectPlatformAccessible:
         pipeline = block_dict['pipeline']
         uuid = block_dict['uuid']
 
-        if self.project_platform_activated:
+        if hydrate_configuration and self.project_platform_activated:
             # self.project_path
             #   /home/src/default_platform/default_repo/dbt/demo
             # node['original_file_path']

mage_ai/data_preparation/models/block/platform/utils.py

@@ -8,12 +8,20 @@ from mage_ai.settings.platform import (
 )
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.files import find_directory
+from mage_ai.shared.path_fixer import get_path_parts, remove_base_repo_directory_name
 
 
-def from_another_project(file_path: str) -> bool:
+def from_another_project(file_path: str, other_file_path: str = None) -> bool:
     if not file_path or not project_platform_activated():
         return False
 
+    if other_file_path:
+        tup1 = get_path_parts(remove_base_repo_directory_name(file_path))
+        tup2 = get_path_parts(remove_base_repo_directory_name(other_file_path))
+
+        if len(tup1) >= 2 and len(tup2) >= 2:
+            return tup1[1] != tup2[1]
+
     active_repo_path = get_repo_path(root_project=False)
     paths = get_repo_paths_for_file_path(
         file_path,

mage_ai/data_preparation/models/file.py

@@ -484,7 +484,7 @@ def traverse(
 
 def __should_update_dbt_cache(dir_path: str, filename: str) -> bool:
     project_model = Project(root_project=True)
-    if project_model and project_model.is_feature_enabled(FeatureUUID.DBT_V2):
+    if project_model and project_model.is_feature_enabled(FeatureUUID.DBT_V2) and dir_path:
         # If the file is a SQL or YAML file
         if (
             filename.endswith('.sql') or

mage_ai/data_preparation/models/global_hooks/models.py

@@ -1085,9 +1085,9 @@ class GlobalHooks(BaseDataClass):
             user=user,
         )
 
-    def save(self, file_path: str = None) -> None:
+    def save(self, file_path: str = None, repo_path: str = None) -> None:
         if not file_path:
-            file_path = self.file_path()
+            file_path = self.file_path(repo_path=repo_path)
 
         content_original = None
         if os.path.exists(file_path):

mage_ai/data_preparation/models/project/__init__.py

@@ -3,11 +3,13 @@ from typing import Dict, List
 
 import aiohttp
 
+from mage_ai.cluster_manager.constants import KUBE_NAMESPACE, ClusterType
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
-from mage_ai.data_preparation.repo_manager import get_repo_config
+from mage_ai.data_preparation.repo_manager import get_cluster_type, get_repo_config
 from mage_ai.server.constants import VERSION
 from mage_ai.settings.platform import (
     active_project_settings,
+    platform_settings,
     project_platform_activated,
     project_platform_settings,
 )
@@ -35,6 +37,34 @@ class Project():
             root_project=self.root_project,
         )
         self.version = VERSION
+        self._features = None
+
+        self.__repo_config_root_project = None
+
+        if project_platform_activated():
+            self.__repo_config_root_project = get_repo_config(
+                repo_path=get_repo_path(root_project=True),
+                root_project=True,
+            )
+
+    @property
+    def workspace_config_defaults(self) -> Dict:
+        config = self.repo_config.workspace_config_defaults or {}
+        cluster_type = get_cluster_type()
+        try:
+            if cluster_type == ClusterType.K8S:
+                from mage_ai.cluster_manager.kubernetes.workload_manager import (
+                    WorkloadManager,
+                )
+                workload_manager = WorkloadManager(os.getenv(KUBE_NAMESPACE))
+                k8s_default_values = workload_manager.get_default_values()
+
+                if not config.get('k8s'):
+                    config['k8s'] = k8s_default_values
+        except Exception:
+            pass
+
+        return config
 
     @property
     def help_improve_mage(self) -> bool:
@@ -50,15 +80,51 @@ class Project():
 
     @property
     def features(self) -> Dict:
+        if self._features:
+            return self._features
+
+        self._features = {}
+        features = self.features_defined
+
+        for uuid in FeatureUUID:
+            key = uuid.value
+            self._features[key] = features.get(key) if features else None
+
+        return self._features
+
+    @features.setter
+    def features(self, x):
+        self._features = x
+
+    @property
+    def features_defined(self) -> Dict:
         data = {}
-        features = self.repo_config.features
+        features = self.repo_config.features if self.repo_config else {}
+
+        if project_platform_activated() and self.__repo_config_root_project:
+            settings = platform_settings()
+            if settings.get('features') and (settings.get('features') or {}).get('override'):
+                features.update(self.features_override)
 
         for uuid in FeatureUUID:
             key = uuid.value
-            data[key] = features.get(key) if features else None
+            if FeatureUUID.PROJECT_PLATFORM == uuid:
+                data[key] = project_platform_activated()
+            else:
+                if features and key in features:
+                    data[key] = features.get(key)
 
         return data
 
+    @property
+    def features_override(self) -> Dict:
+        if project_platform_activated() and self.__repo_config_root_project:
+            settings = self.platform_settings()
+            if settings.get('features') and (settings.get('features') or {}).get('override'):
+                return self.__repo_config_root_project.features or {}
+
+        return {}
+
     @property
     def emr_config(self) -> Dict:
         return self.repo_config.emr_config or None
@@ -85,6 +151,10 @@ class Project():
 
         return False
 
+    def platform_settings(self) -> Dict:
+        if project_platform_activated():
+            return platform_settings(mage_projects_only=True)
+
     def repo_path_for_database_query(self, key: str) -> List[str]:
         if self.settings:
             query_arr = dig(self.settings, ['database', 'query', key])
@@ -102,7 +172,7 @@ class Project():
         return project_platform_settings(mage_projects_only=True)
 
     def is_feature_enabled(self, feature_name: FeatureUUID) -> bool:
-        feature_enabled = self.repo_config.features.get(feature_name.value, False)
+        feature_enabled = self.features.get(feature_name.value, False)
 
         if is_debug():
             print(f'[Project.is_feature_enabled]: {feature_name} | {feature_enabled}')

mage_ai/data_preparation/models/project/constants.py

@@ -12,3 +12,4 @@ class FeatureUUID(str, Enum):
     LOCAL_TIMEZONE = 'display_local_timezone'
     NOTEBOOK_BLOCK_OUTPUT_SPLIT_VIEW = 'notebook_block_output_split_view'
     OPERATION_HISTORY = 'operation_history'
+    PROJECT_PLATFORM = 'project_platform'

mage_ai/data_preparation/repo_manager.py

@@ -64,6 +64,7 @@ class RepoConfig:
         self.s3_bucket = None
         self.s3_path_prefix = None
         self.logging_config = None
+        self.variables_dir = None
         self.variables_retention_period = None
 
         from mage_ai.data_preparation.shared.utils import get_template_vars
@@ -122,6 +123,7 @@ class RepoConfig:
             self.openai_api_key = repo_config.get('openai_api_key')
             self.pipelines = repo_config.get('pipelines')
             self.retry_config = repo_config.get('retry_config')
+            self.workspace_config_defaults = repo_config.get('workspace_config_defaults')
 
             self.ldap_config = repo_config.get('ldap_config')
 
@@ -194,6 +196,7 @@ class RepoConfig:
             spark_config=self.spark_config,
             variables_dir=self.remote_variables_dir if remote else self.variables_dir,
             variables_retention_period=self.variables_retention_period,
+            workspace_config_defaults=self.workspace_config_defaults
         )
 
     def save(self, **kwargs) -> None: