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

mage_ai/data_preparation/models/custom_templates/custom_pipeline_template.py

@@ -1,13 +1,18 @@
 import os
 import shutil
-import yaml
 from dataclasses import dataclass, field
-from mage_ai.data_preparation.models.pipeline import Pipeline
+from typing import Dict, List
+
+import yaml
+
 from mage_ai.data_preparation.models.custom_templates.constants import (
     DIRECTORY_FOR_PIPELINE_TEMPLATES,
     METADATA_FILENAME_WITH_EXTENSION,
 )
-from mage_ai.data_preparation.models.custom_templates.utils import custom_templates_directory
+from mage_ai.data_preparation.models.custom_templates.utils import (
+    custom_templates_directory,
+)
+from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import (
     TRIGGER_FILE_NAME,
     Trigger,
@@ -16,12 +21,10 @@ from mage_ai.data_preparation.models.triggers import (
     load_trigger_configs,
 )
 from mage_ai.orchestration.db.models.schedules import PipelineSchedule
-from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.config import BaseConfig
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.io import safe_write
 from mage_ai.shared.utils import clean_name
-from typing import Dict, List
 
 
 @dataclass
@@ -29,12 +32,13 @@ class CustomPipelineTemplate(BaseConfig):
     description: str = None
     name: str = None
     pipeline: Dict = field(default_factory=dict)
+    repo_path: str = None
     tags: List = field(default_factory=list)
     template_uuid: str = None
     user: Dict = field(default_factory=dict)
 
     @classmethod
-    def load(self, template_uuid: str = None, uuid: str = None):
+    def load(self, repo_path: str, template_uuid: str = None, uuid: str = None):
         uuid_use = uuid
         template_uuid_use = template_uuid
 
@@ -50,12 +54,13 @@ class CustomPipelineTemplate(BaseConfig):
 
         try:
             config_path_metadata = os.path.join(
-                get_repo_path(),
+                repo_path,
                 uuid_use,
                 METADATA_FILENAME_WITH_EXTENSION,
             )
             custom_template = super().load(config_path_metadata)
             custom_template.template_uuid = template_uuid_use
+            custom_template.repo_path = repo_path
 
             return custom_template
         except Exception as err:
@@ -78,6 +83,7 @@ class CustomPipelineTemplate(BaseConfig):
             description=description,
             name=name,
             pipeline=pipeline_dict,
+            repo_path=pipeline.repo_path,
             template_uuid=clean_name(template_uuid, [os.sep]) if template_uuid else template_uuid,
         )
 
@@ -118,7 +124,7 @@ class CustomPipelineTemplate(BaseConfig):
     @property
     def metadata_file_path(self) -> str:
         return os.path.join(
-            get_repo_path(),
+            self.repo_path,
             self.uuid,
             METADATA_FILENAME_WITH_EXTENSION,
         )
@@ -126,16 +132,24 @@ class CustomPipelineTemplate(BaseConfig):
     @property
     def triggers_file_path(self) -> str:
         return os.path.join(
-            get_repo_path(),
+            self.repo_path,
             self.uuid,
             TRIGGER_FILE_NAME,
         )
 
     def build_pipeline(self) -> Pipeline:
-        return Pipeline(clean_name(self.template_uuid), config=self.pipeline)
+        return Pipeline(
+            clean_name(self.template_uuid),
+            config=self.pipeline,
+            repo_path=self.repo_path,
+        )
 
     def create_pipeline(self, name: str) -> Pipeline:
-        pipeline = Pipeline(clean_name(name), config=self.pipeline)
+        pipeline = Pipeline(
+            clean_name(name),
+            repo_path=self.repo_path,
+            config=self.pipeline,
+        )
         os.makedirs(os.path.dirname(pipeline.config_path), exist_ok=True)
         pipeline.save()
 
@@ -151,10 +165,13 @@ class CustomPipelineTemplate(BaseConfig):
         return pipeline
 
     def to_dict(self) -> Dict:
-        return merge_dict(self.to_dict_base(), dict(
-            template_uuid=self.template_uuid,
-            uuid=self.uuid,
-        ))
+        return merge_dict(
+            self.to_dict_base(),
+            dict(
+                template_uuid=self.template_uuid,
+                uuid=self.uuid,
+            ),
+        )
 
     def to_dict_base(self) -> Dict:
         return dict(
@@ -178,4 +195,4 @@ class CustomPipelineTemplate(BaseConfig):
         safe_write(file_path, content)
 
     def delete(self) -> None:
-        shutil.rmtree(os.path.join(get_repo_path(), self.uuid))
+        shutil.rmtree(os.path.join(self.repo_path, self.uuid))

mage_ai/data_preparation/models/custom_templates/utils.py

@@ -71,7 +71,7 @@ def group_and_hydrate_files(
     custom_templates = []
 
     for template_uuid, _ in groups.items():
-        custom_template = custom_template_class.load(template_uuid=template_uuid)
+        custom_template = custom_template_class.load(get_repo_path(), template_uuid=template_uuid)
         if custom_template:
             custom_templates.append(custom_template)
 

mage_ai/data_preparation/models/file.py

@@ -145,6 +145,7 @@ class File:
         file_selector = None
 
         if exclude_pattern is not None or pattern is not None:
+
             def __select(x: Dict, check_file_path=check_file_path, pattern=pattern):
                 filename = x.get('path') if check_file_path else x.get('name')
                 checks = []
@@ -157,6 +158,7 @@ class File:
             file_selector = __select
 
         if exclude_dir_pattern is not None:
+
             def __select(x: Dict, pattern=pattern):
                 filename = x.get('name')
                 checks = []
@@ -169,6 +171,7 @@ class File:
         pipeline_count_mapping = None
         if include_pipeline_count:
             from mage_ai.cache.block import BlockCache
+
             pipeline_count_mapping = BlockCache().get_pipeline_count_mapping()
 
         return traverse(
@@ -200,7 +203,7 @@ class File:
         if dir_path.startswith(PIPELINES_FOLDER_PREFIX) and filename == 'triggers.yaml':
             from mage_ai.data_preparation.models.triggers import load_trigger_configs
 
-            pipeline_uuid = dir_path[len(PIPELINES_FOLDER_PREFIX):]
+            pipeline_uuid = dir_path[len(PIPELINES_FOLDER_PREFIX) :]
             load_trigger_configs(content, pipeline_uuid=pipeline_uuid, raise_exception=True)
 
     @classmethod
@@ -239,8 +242,9 @@ class File:
             number_of_file_versions = len(file_versions)
 
             if number_of_file_versions >= MAX_NUMBER_OF_FILE_VERSIONS:
-                number_of_file_versions_to_delete = 1 + \
-                    (number_of_file_versions - MAX_NUMBER_OF_FILE_VERSIONS)
+                number_of_file_versions_to_delete = 1 + (
+                    number_of_file_versions - MAX_NUMBER_OF_FILE_VERSIONS
+                )
                 for fn in sorted(file_versions)[:number_of_file_versions_to_delete]:
                     fn_path = os.path.join(file_path_versions_dir, fn)
                     os.remove(fn_path)
@@ -339,8 +343,9 @@ class File:
                 file_content = fp.read()
             return file_content
         except FileNotFoundError as err:
-            print('file.content')
-            print(err)
+            print(f'[ERROR] File.content: {err}')
+            if is_debug():
+                raise err
         return ''
 
     async def content_async(self):
@@ -370,11 +375,14 @@ class File:
             os.path.join(self.repo_path, ''),
             '',
         )
-        return [File(
-            v,
-            file_path_versions_dir_without_repo,
-            self.repo_path,
-        ) for v in sorted(file_versions, reverse=True)]
+        return [
+            File(
+                v,
+                file_path_versions_dir_without_repo,
+                self.repo_path,
+            )
+            for v in sorted(file_versions, reverse=True)
+        ]
 
     def update_content(self, content: str):
         self.write(
@@ -464,8 +472,7 @@ def ensure_file_is_in_project(file_path: str) -> None:
     full_file_path = get_absolute_path(file_path)
     full_repo_path = get_absolute_path(get_repo_path(file_path=file_path))
     if full_repo_path != os.path.commonpath([full_file_path, full_repo_path]):
-        raise FileNotInProjectError(
-            f'File at path: {file_path} is not in the project directory.')
+        raise FileNotInProjectError(f'File at path: {file_path} is not in the project directory.')
 
 
 def traverse(
@@ -508,17 +515,25 @@ def traverse(
 
         entry_path = entry.path
         if entry.is_dir(follow_symlinks=False) or os.path.isdir(entry_path):
-            return True if dir_selector is None else dir_selector(dict(
-                depth=depth + 1,
-                name=str(entry.name),
-                path=entry.path,
-            ))
+            return (
+                True
+                if dir_selector is None
+                else dir_selector(
+                    dict(
+                        depth=depth + 1,
+                        name=str(entry.name),
+                        path=entry.path,
+                    )
+                )
+            )
         else:
-            return file_selector(dict(
-                depth=depth + 1,
-                name=str(entry.name),
-                path=entry.path,
-            ))
+            return file_selector(
+                dict(
+                    depth=depth + 1,
+                    name=str(entry.name),
+                    path=entry.path,
+                )
+            )
 
         return True
 
@@ -533,7 +548,8 @@ def traverse(
             file_selector=file_selector,
             include_pipeline_count=include_pipeline_count,
             pipeline_count_mapping=pipeline_count_mapping,
-        ) for entry in sorted(
+        )
+        for entry in sorted(
             filter(__filter, os.scandir(path)),
             key=lambda entry: entry.name,
         )
@@ -546,11 +562,7 @@ 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) and dir_path:
         # If the file is a SQL or YAML file
-        if (
-            filename.endswith('.sql') or
-            filename.endswith('.yml') or
-            filename.endswith('.yaml')
-        ):
+        if filename.endswith('.sql') or filename.endswith('.yml') or filename.endswith('.yaml'):
             base_dir_path = Path(dir_path).parts[0].lower()
             # If the file doesn’t exist in a block’s folder (e.g. data_loaders/, pipelines/)
             if base_dir_path not in IGNORE_DIRECTORY_NAMES:
@@ -595,4 +607,5 @@ def update_file_cache() -> None:
     project_model = Project(root_project=True)
     if project_model and project_model.is_feature_enabled(FeatureUUID.COMMAND_CENTER):
         from mage_ai.cache.file import FileCache
+
         FileCache().invalidate()

mage_ai/data_preparation/models/global_data_product/__init__.py

@@ -1,21 +1,22 @@
 import os
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import yaml
 from dateutil.relativedelta import relativedelta
 
-from mage_ai.data_preparation.models.global_data_product.constants import (
+from mage_ai.data_preparation.models.block.settings.global_data_products.models import (
     GlobalDataProductObjectType,
 )
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
-from mage_ai.settings.repo import get_repo_path
+from mage_ai.settings.platform.utils import full_paths_for_all_projects
 from mage_ai.shared.array import find
 from mage_ai.shared.dates import week_of_month
 from mage_ai.shared.hash import extract, index_by
 from mage_ai.shared.io import safe_write
+from mage_ai.shared.path_fixer import add_absolute_path, get_path_parts
 from mage_ai.shared.utils import clean_name
 
 
@@ -26,6 +27,7 @@ class GlobalDataProduct:
     object_uuid: str = None
     outdated_after: Dict = field(default_factory=dict)
     outdated_starting_at: Dict = field(default_factory=dict)
+    repo_path: str = None
     settings: Dict = field(default_factory=dict)
 
     def __init__(self, uuid: str, **kwargs):
@@ -33,37 +35,50 @@ class GlobalDataProduct:
         self.object_uuid = kwargs.get('object_uuid')
         self.outdated_after = kwargs.get('outdated_after')
         self.outdated_starting_at = kwargs.get('outdated_starting_at')
+        self.repo_path = kwargs.get('repo_path')
         self.settings = kwargs.get('settings')
         self.uuid = clean_name(uuid) if uuid else uuid
         self._object = None
 
     @classmethod
-    def file_path(self) -> str:
-        return os.path.join(get_repo_path(), 'global_data_products.yaml')
+    def file_path(cls, repo_path: str) -> str:
+        return os.path.join(repo_path, 'global_data_products.yaml')
 
     @classmethod
-    def load_all(self, file_path: str = None) -> List['GlobalDataProduct']:
-        file_path_to_use = file_path or self.file_path()
-        arr = []
+    def load_all(cls, repo_path: Optional[str] = None) -> List['GlobalDataProduct']:
+        paths = []
+        if repo_path:
+            paths.append(repo_path)
+        else:
+            paths += full_paths_for_all_projects()
 
-        if not os.path.exists(file_path_to_use):
-            return []
+        arr = []
 
-        with open(file_path_to_use) as fp:
-            content = fp.read()
-            if content:
-                yaml_config = yaml.safe_load(content) or {}
-                for uuid, config_attributes in yaml_config.items():
-                    arr.append(self(
-                        uuid,
-                        **config_attributes,
-                    ))
+        for file_path in paths:
+            file_path_to_use = cls.file_path(file_path)
+            if not os.path.exists(file_path_to_use):
+                continue
+
+            with open(file_path_to_use) as fp:
+                content = fp.read()
+                if content:
+                    yaml_config = yaml.safe_load(content) or {}
+                    for uuid, config_attributes in yaml_config.items():
+                        arr.append(
+                            cls(
+                                uuid,
+                                **config_attributes,
+                            )
+                        )
 
         return arr
 
     @classmethod
-    def get(self, uuid: str, file_path: str = None):
-        return find(lambda x: x.uuid == uuid, self.load_all(file_path=file_path))
+    def get(cls, uuid: str, repo_path: Optional[str] = None, project: Optional[str] = None):
+        if project:
+            repo_path = add_absolute_path(project)
+
+        return find(lambda x: x.uuid == uuid, cls.load_all(repo_path))
 
     @property
     def pipeline(self) -> Pipeline:
@@ -71,7 +86,7 @@ class GlobalDataProduct:
             if self._object:
                 return self._object
             else:
-                self._object = Pipeline.get(self.object_uuid)
+                self._object = Pipeline.get(self.object_uuid, repo_path=self.repo_path)
 
         return self._object
 
@@ -137,12 +152,15 @@ class GlobalDataProduct:
 
     def get_outdated_at_delta(self, in_seconds: bool = False) -> relativedelta:
         outdated_after = self.outdated_after or {}
-        delta = extract(outdated_after, [
-            'months',
-            'seconds',
-            'weeks',
-            'years',
-        ])
+        delta = extract(
+            outdated_after,
+            [
+                'months',
+                'seconds',
+                'weeks',
+                'years',
+            ],
+        )
 
         if len(delta) >= 1:
             d = relativedelta(**delta)
@@ -232,8 +250,7 @@ class GlobalDataProduct:
         status: PipelineRun.PipelineRunStatus = None,
     ) -> List[PipelineRun]:
         pipeline_runs = (
-            PipelineRun.
-            select(
+            PipelineRun.select(
                 PipelineRun.backfill_id,
                 PipelineRun.completed_at,
                 PipelineRun.created_at,
@@ -249,20 +266,20 @@ class GlobalDataProduct:
                 PipelineRun.updated_at,
                 PipelineRun.variables,
                 PipelineSchedule.global_data_product_uuid,
-            ).
-            join(PipelineSchedule, PipelineRun.pipeline_schedule_id == PipelineSchedule.id).
-            filter(
+            )
+            .join(
+                PipelineSchedule,
+                PipelineRun.pipeline_schedule_id == PipelineSchedule.id,
+            )
+            .filter(
                 PipelineRun.pipeline_uuid == self.object_uuid,
                 PipelineSchedule.global_data_product_uuid == self.uuid,
             )
         )
 
         if status:
-            pipeline_runs = (
-                pipeline_runs.
-                filter(
-                    PipelineRun.status == status,
-                )
+            pipeline_runs = pipeline_runs.filter(
+                PipelineRun.status == status,
             )
 
         pipeline_runs = pipeline_runs.order_by(PipelineRun.execution_date.desc())
@@ -270,23 +287,33 @@ class GlobalDataProduct:
         if limit is not None:
             pipeline_runs = pipeline_runs.limit(limit)
 
-        return [PipelineRun(
-            backfill_id=row.backfill_id,
-            completed_at=row.completed_at,
-            created_at=row.created_at,
-            event_variables=row.event_variables,
-            execution_date=row.execution_date,
-            executor_type=row.executor_type,
-            global_data_product_uuid=row.global_data_product_uuid,
-            id=row.id,
-            metrics=row.metrics,
-            passed_sla=row.passed_sla,
-            pipeline_schedule_id=row.pipeline_schedule_id,
-            pipeline_uuid=row.pipeline_uuid,
-            status=row.status,
-            updated_at=row.updated_at,
-            variables=row.variables,
-        ) for row in pipeline_runs.all()]
+        return [
+            PipelineRun(
+                backfill_id=row.backfill_id,
+                completed_at=row.completed_at,
+                created_at=row.created_at,
+                event_variables=row.event_variables,
+                execution_date=row.execution_date,
+                executor_type=row.executor_type,
+                global_data_product_uuid=row.global_data_product_uuid,
+                id=row.id,
+                metrics=row.metrics,
+                passed_sla=row.passed_sla,
+                pipeline_schedule_id=row.pipeline_schedule_id,
+                pipeline_uuid=row.pipeline_uuid,
+                status=row.status,
+                updated_at=row.updated_at,
+                variables=row.variables,
+            )
+            for row in pipeline_runs.all()
+        ]
+
+    @property
+    def project_name(self) -> Optional[str]:
+        if self.repo_path:
+            parts = get_path_parts(self.repo_path)
+            if parts and len(parts) >= 3:
+                return parts[1]
 
     def to_dict(self, include_uuid: bool = False) -> Dict:
         data = dict(
@@ -294,6 +321,8 @@ class GlobalDataProduct:
             object_uuid=self.object_uuid,
             outdated_after=self.outdated_after,
             outdated_starting_at=self.outdated_starting_at,
+            project=self.project_name,
+            repo_path=self.repo_path,
             settings=self.settings,
         )
 
@@ -305,8 +334,8 @@ class GlobalDataProduct:
     def delete(self) -> None:
         self.save(delete=True)
 
-    def save(self, delete: bool = False, file_path: str = None) -> None:
-        mapping = index_by(lambda x: x.uuid, self.load_all())
+    def save(self, delete: bool = False) -> None:
+        mapping = index_by(lambda x: x.uuid, self.load_all(self.repo_path))
 
         if delete:
             mapping.pop(self.uuid, None)
@@ -324,7 +353,7 @@ class GlobalDataProduct:
                     mapping_new[uuid][k] = v
 
         content = yaml.safe_dump(mapping_new)
-        safe_write(file_path or self.file_path(), content)
+        safe_write(self.file_path(self.repo_path), content)
 
     def update(self, payload: Dict) -> None:
         for key in [
@@ -332,6 +361,7 @@ class GlobalDataProduct:
             'object_uuid',
             'outdated_after',
             'outdated_starting_at',
+            'repo_path',
             'settings',
         ]:
             value = payload.get(key)

mage_ai/data_preparation/models/global_hooks/models.py

@@ -421,6 +421,7 @@ class Hook(BaseDataClass):
             if with_trigger or (self.run_settings and self.run_settings.with_trigger):
                 pipeline_run = trigger_pipeline(
                     self.pipeline.uuid,
+                    repo_path=self.pipeline.repo_path,
                     variables=variables,
                     check_status=check_status and not asynchronous,
                     error_on_failure=error_on_failure,

mage_ai/data_preparation/models/interfaces.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional, Sequence, Union
+
+
+@dataclass
+class PipelineInterface:
+    get_block: Callable[[str], 'BlockInterface']
+    uuid: str
+    variable_manager: Any
+
+
+@dataclass
+class BlockInterface:
+    configuration: Dict[str, bool]
+    get_data_integration_settings: Callable[..., Dict[str, Optional[Dict[str, Union[str, Dict]]]]]
+    get_spark_session: Callable[..., Any]
+    get_spark_session_from_global_vars: Callable[..., Any]
+    pipeline: PipelineInterface
+    upstream_blocks: Sequence['BlockInterface']
+    uuid: str
+    variable_manager: Any
+
+
+@dataclass
+class VariableInterface:
+    is_partial_data_readable: Callable[[], bool]
+    items_count: Callable[..., Optional[int]]
+    read_partial_data: Callable
+    uuid: str