mage-ai 0.9.58__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/shared/environments.py

@@ -4,6 +4,10 @@ import sys
 from mage_ai.shared.constants import ENV_DEV, ENV_PROD, ENV_STAGING, ENV_TEST
 
 
+def is_deus_ex_machina():
+    return int(os.getenv('DEUS_EX_MACHINA', 0) or 0) == 1
+
+
 def is_debug():
     return int(os.getenv('DEBUG', 0) or 0) == 1
 

mage_ai/shared/files.py

@@ -1,5 +1,7 @@
+import glob
 import os
-from typing import Callable, List
+from pathlib import Path
+from typing import Callable, List, Tuple
 
 
 def reverse_readline(filename, buf_size=8192):
@@ -63,3 +65,52 @@ def find_directory(top_level_path: str, comparator: Callable) -> str:
             full_path = os.path.join(path, name)
             if comparator(full_path):
                 return full_path
+
+
+def get_absolute_paths_from_all_files(
+    starting_full_path_directory: str,
+    comparator: Callable = None,
+    include_hidden_files: bool = False,
+    parse_values: Callable = None,
+) -> List[Tuple[str, int, str]]:
+    dir_path = os.path.join(starting_full_path_directory, './**/*')
+
+    arr = []
+    for filename in glob.iglob(dir_path, recursive=True):
+        absolute_path = os.path.abspath(filename)
+
+        if os.path.isfile(absolute_path) and \
+                (not include_hidden_files or not absolute_path.startswith('.')) and \
+                (not comparator or comparator(absolute_path)):
+
+            value = (absolute_path, os.path.getsize(filename), round(os.path.getmtime(filename)))
+            arr.append(parse_values(value) if parse_values else value)
+
+    return arr
+
+
+def find_file_from_another_file_path(file_path: str, comparator) -> str:
+    if not file_path or not comparator:
+        return
+
+    if not os.path.isdir(file_path):
+        file_path = os.path.dirname(file_path)
+
+    parts = Path(file_path).parts
+
+    absolute_file_path = None
+
+    while len(parts) > 1 and absolute_file_path is None:
+        if len(parts) == 0:
+            return
+
+        fp = os.path.join(*parts)
+        for fn in os.listdir(fp):
+            afp = os.path.join(fp, fn)
+            if afp is not None and comparator(afp):
+                absolute_file_path = afp
+                break
+
+        parts = parts[:-1]
+
+    return absolute_file_path

mage_ai/shared/models.py

@@ -1,3 +1,4 @@
+import inspect
 import typing
 from dataclasses import dataclass, make_dataclass
 from enum import Enum
@@ -119,7 +120,7 @@ class BaseClass:
             else:
                 return value
 
-        is_data_class = issubclass(annotation, BaseDataClass)
+        is_data_class = inspect.isclass(annotation) and issubclass(annotation, BaseDataClass)
         if is_data_class:
             if is_dict_class:
                 if ignore_empty and len(value) == 0:

mage_ai/shared/parsers.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from enum import Enum
 from json import JSONDecoder
 
 import numpy as np
@@ -23,8 +24,14 @@ MAX_ITEMS_IN_SAMPLE_OUTPUT = 20
 
 
 def encode_complex(obj):
-    if isinstance(obj, BaseModel):
+    if isinstance(obj, set):
+        return list(obj)
+    elif isinstance(obj, BaseModel):
         return obj.__class__.__name__
+    elif obj.__class__.__name__ == 'BaseDataClass':
+        return obj.to_dict()
+    elif isinstance(obj, Enum):
+        return obj.value
     elif hasattr(obj, 'isoformat') and 'method' in type(obj.isoformat).__name__:
         return obj.isoformat()
     elif isinstance(obj, np.integer):

mage_ai/shared/path_fixer.py

@@ -5,6 +5,39 @@ from typing import Tuple
 from mage_ai.settings.platform import get_repo_paths_for_file_path
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.settings.utils import base_repo_dirname, base_repo_name, base_repo_path
+from mage_ai.shared.files import find_directory
+
+
+def add_absolute_path(file_path: str, add_base_repo_path: bool = True) -> str:
+    repo_path = get_repo_path(root_project=True)
+    repo_path_name = Path(repo_path).name
+
+    full_path = None
+    if os.path.isabs(file_path):
+        full_path = file_path
+    else:
+        try:
+            # If the relative path already contains the root project name, return it
+            diff = Path(file_path).relative_to(repo_path_name)
+            relative_file_path = str(diff)
+            full_path = os.path.join(repo_path, relative_file_path)
+        except ValueError:
+            # Relative path doesn’t contain the repo name
+            # Check to see if the first directory is exists right under the root folder
+            parts = Path(file_path).parts
+            if len(parts) >= 1:
+                full_path = os.path.join(repo_path, parts[0])
+                if os.path.exists(full_path):
+                    full_path = os.path.join(repo_path, file_path)
+                else:
+                    full_path = find_directory(repo_path, lambda x: x.endswith(file_path))
+
+    if full_path:
+        # Remove the repo path
+        if add_base_repo_path:
+            return str(full_path)
+        else:
+            return str(Path(full_path).relative_to(base_repo_path()))
 
 
 def add_root_repo_path_to_relative_path(relative_file_path: str) -> str:
@@ -55,6 +88,21 @@ def remove_base_repo_name(file_path: str) -> str:
         return file_path
 
 
+def remove_base_repo_path_or_name(file_path: str) -> str:
+    if not file_path:
+        return
+
+    if os.path.isabs(file_path):
+        value = str(remove_base_repo_path(file_path))
+    else:
+        value = str(remove_base_repo_name(file_path))
+
+    if len(value) == 0 or value == '.':
+        return base_repo_name()
+    else:
+        return value
+
+
 def remove_base_repo_directory_name(file_path: str) -> str:
     """
     Removes /home/src from
@@ -90,18 +138,27 @@ def get_path_parts(file_path: str) -> Tuple[str, str, str]:
     full_path_relative = paths['full_path_relative']
     path = paths['path']
 
+    paths_to_remove = []
     file_path_base = None
 
+    if os.path.isabs(file_path):
+        paths_to_remove = [
+            paths['full_path'],
+            path,
+        ]
+    else:
+        paths_to_remove = [
+            full_path_relative,
+            path,
+        ]
+
     # The file path either has both the root project name and the nested project name
     # or just has the nested project name in it.
-    for path2 in [
-        full_path_relative,
-        path,
-    ]:
+    for path_inner in paths_to_remove:
         if file_path_base:
             break
         try:
-            file_path_base = Path(file_path).relative_to(path2)
+            file_path_base = Path(file_path).relative_to(path_inner)
         except ValueError:
             pass
 

mage_ai/shared/retry.py

@@ -9,6 +9,7 @@ def retry(
     exponential_backoff: bool = True,
     logger=None,
     logging_tags: Dict = None,
+    retry_metadata: Dict = None,
 ):
     """
     Create the decorator to retry a method
@@ -31,6 +32,9 @@ def retry(
             curr_delay = delay
             while attempt <= max_attempts:
                 try:
+                    if retry_metadata:
+                        retry_metadata['attempts'] = (retry_metadata.get('attempts', 0) or 0) + 1
+
                     return func(*args, **kwargs)
                 except Exception as e:
                     if logger is None:

mage_ai/shared/strings.py

@@ -52,3 +52,7 @@ def remove_extension_from_filename(filename: str) -> str:
 
 def singularize(word: str) -> str:
     return inflection.singularize(word)
+
+
+def capitalize_remove_underscore_lower(word: str) -> str:
+    return word.replace('_', ' ').lower().capitalize()

mage_ai/tests/api/endpoints/test_blocks.py

@@ -1,4 +1,6 @@
-from mage_ai.data_preparation.models.constants import BlockType
+import os
+
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
@@ -8,6 +10,7 @@ from mage_ai.tests.api.endpoints.mixins import (
     build_list_endpoint_tests,
     build_update_endpoint_tests,
 )
+from mage_ai.tests.shared.mixins import DBTMixin
 
 
 def get_pipeline(self):
@@ -105,3 +108,100 @@ build_delete_endpoint_tests(
     resource='block',
     resource_parent='pipelines',
 )
+
+
+class DBTBlockAPIEndpointTest(BaseAPIEndpointTest, DBTMixin):
+    pass
+
+
+def __assert_after_create(self):
+    pipeline = get_pipeline(self)
+    blocks = pipeline.blocks_by_uuid.values()
+    blocks_dbt = list(filter(
+        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.SQL,
+        blocks,
+    ))
+
+    self.assertEqual(len(blocks), 5)
+    self.assertEqual(len(blocks_dbt), 1)
+
+    block = blocks_dbt[0]
+    self.assertEqual({
+      'dbt_project_name': 'dbt/test_project',
+      'file_path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
+      'file_source': {
+        'path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
+        'project_path': 'dbt/test_project',
+      },
+    }, block.configuration)
+
+    return True
+
+
+def __assert_after_create_yaml(self):
+    pipeline = get_pipeline(self)
+    blocks = pipeline.blocks_by_uuid.values()
+    blocks_dbt = list(filter(
+        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.YAML,
+        blocks,
+    ))
+
+    self.assertEqual(len(blocks), 5)
+    self.assertEqual(len(blocks_dbt), 1)
+
+    block = blocks_dbt[0]
+    self.assertEqual({
+      'dbt_project_name': 'dbt/test_project',
+      'file_source': {
+        'path': 'dbts/dbt_yaml.yaml',
+        'project_path': 'dbt/test_project',
+      },
+    }, block.configuration)
+
+    return True
+
+
+build_create_endpoint_tests(
+    DBTBlockAPIEndpointTest,
+    assert_after_create_count=__assert_after_create,
+    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
+    build_payload=lambda self: dict(
+        name=self.faker.unique.name(),
+        type=BlockType.DBT,
+        language=BlockLanguage.SQL,
+        configuration=dict(
+            file_path=os.path.join(
+                os.path.basename(self.dbt_directory),
+                'test_project',
+                'models',
+                'example',
+                'my_first_dbt_model.sql',
+            ),
+        ),
+    ),
+    get_resource_parent_id=lambda self: self.pipeline.uuid,
+    resource='block',
+    resource_parent='pipelines',
+)
+
+
+build_create_endpoint_tests(
+    DBTBlockAPIEndpointTest,
+    test_uuid='dbt_yaml',
+    assert_after_create_count=__assert_after_create_yaml,
+    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
+    build_payload=lambda self: dict(
+        name='dbt_yaml',
+        type=BlockType.DBT,
+        language=BlockLanguage.YAML,
+        configuration=dict(
+            dbt_project_name=os.path.join(
+                os.path.basename(self.dbt_directory),
+                'test_project',
+            ),
+        ),
+    ),
+    get_resource_parent_id=lambda self: self.pipeline.uuid,
+    resource='block',
+    resource_parent='pipelines',
+)

mage_ai/tests/api/operations/base/test_base_with_user_permissions.py

@@ -14,7 +14,10 @@ class BaseOperationsWithUserAuthenticationAndPermissionsTest(Base):
     def setUp(self):
         self.set_up()
 
-        self.user.save()
+        try:
+            self.user.save()
+        except Exception as err:
+            print(err)
 
         role = Role.create(name=self.faker.unique.name())
         permission = Permission.create(

mage_ai/tests/base_test.py

@@ -26,7 +26,18 @@ else:
         @classmethod
         def setUpClass(self):
             super().setUpClass()
-            self.repo_path = os.path.join(os.getcwd(), 'test')
+
+            repo_path = None
+            for fp in [
+                os.path.dirname(os.getcwd()),
+                os.getcwd(),
+                'test'
+            ]:
+                repo_path = os.path.join(repo_path or '', fp)
+                if not os.path.exists(repo_path):
+                    os.makedirs(repo_path, exist_ok=True)
+
+            self.repo_path = repo_path
             set_repo_path(self.repo_path)
             if not Path(self.repo_path).exists():
                 Path(self.repo_path).mkdir()
@@ -89,7 +100,15 @@ class TestCase(unittest.TestCase):
     @classmethod
     def setUpClass(self):
         super().setUpClass()
-        self.repo_path = os.path.join(os.getcwd(), 'test')
+        repo_path = os.path.join(os.getcwd())
+        if not os.path.exists(repo_path):
+            os.makedirs(repo_path, exist_ok=True)
+
+        repo_path = os.path.join(repo_path, 'test')
+        if not os.path.exists(repo_path):
+            os.makedirs(repo_path, exist_ok=True)
+
+        self.repo_path = repo_path
         set_repo_path(self.repo_path)
         if not Path(self.repo_path).exists():
             Path(self.repo_path).mkdir()

mage_ai/tests/cache/test_pipeline_cache.py

@@ -2,7 +2,7 @@ import json
 import os
 import shutil
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from unittest.mock import patch
 
 from freezegun import freeze_time
@@ -33,7 +33,7 @@ class PipelineCacheTest(BaseApiTestCase):
             ),
         )
 
-        self.cache = CustomPipelineCache()
+        self.cache = CustomPipelineCache(root_project=False)
         with open(os.path.join(CURRENT_FILE_PATH, 'example_pipeline.json')) as f:
             self.pipeline_json = json.loads(f.read())['pipeline']
 
@@ -89,7 +89,7 @@ class PipelineCacheTest(BaseApiTestCase):
         )
 
     def test_update_models(self):
-        now = datetime(3000, 1, 1)
+        now = datetime(2024, 1, 1)
         model = self.pipeline_json.copy()
         model['name'] = self.uuid
 
@@ -106,7 +106,7 @@ class PipelineCacheTest(BaseApiTestCase):
         )
 
     def test_update_model(self):
-        now = datetime(3000, 1, 1)
+        now = datetime(2024, 1, 1)
         model = self.pipeline_json.copy()
         model['name'] = self.uuid
 
@@ -123,7 +123,7 @@ class PipelineCacheTest(BaseApiTestCase):
         )
 
     def test_add_model(self):
-        now = datetime(3000, 1, 1)
+        now = datetime(2024, 1, 1, tzinfo=timezone.utc)
         model = self.pipeline_json.copy()
         model['name'] = self.uuid
 
@@ -146,9 +146,10 @@ class PipelineCacheTest(BaseApiTestCase):
         )
 
     async def test_initialize_cache_for_models(self):
-        now = datetime(3000, 1, 1)
+        now = datetime(2024, 1, 1, tzinfo=timezone.utc)
 
         pipeline_dict = dict(
+            created_at=None,
             description=uuid.uuid4().hex,
             blocks=[
                 dict(
@@ -171,10 +172,12 @@ class PipelineCacheTest(BaseApiTestCase):
             uuid=uuid.uuid4().hex,
         )
 
+        repo_path = self.repo_path
+
         class TestCustomPipeline:
             @classmethod
-            def get_all_pipelines(self, *args, **kwargs):
-                return [pipeline_dict['uuid']]
+            def get_all_pipelines_all_projects(cls, repo_path=repo_path, *args, **kwargs):
+                return [(pipeline_dict['uuid'], repo_path)]
 
             @classmethod
             async def load_metadata(self, pipeline_uuid, **kwargs):

mage_ai/tests/data_preparation/executors/test_block_executor.py

@@ -15,6 +15,7 @@ from mage_ai.data_preparation.models.block.hook.block import HookBlock
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun
+from mage_ai.shared.hash import ignore_keys, merge_dict
 from mage_ai.tests.api.operations.test_base import BaseApiTestCase
 from mage_ai.tests.factory import create_pipeline_with_blocks
 from mage_ai.tests.shared.mixins import build_hooks
@@ -33,6 +34,7 @@ class BlockExecutorTest(BaseApiTestCase):
         self.block.configuration = {}
         self.block.upstream_blocks = []
         self.block.type = BlockType.DBT
+        self.block.replicated_block = None
 
         self.pipeline.get_block.return_value = self.block
         self.pipeline.repo_config.retry_config = {'retries': 3, 'delay': 1}
@@ -75,7 +77,9 @@ class BlockExecutorTest(BaseApiTestCase):
     def test_execute(self):
         analyze_outputs = False
         callback_url = 'http://example.com/callback'
-        global_vars = {'var1': 'value1'}
+        global_vars = {
+            'var1': 'value1',
+        }
         update_status = True
 
         def on_complete(block_uuid, metrics: Dict = None):
@@ -133,7 +137,7 @@ class BlockExecutorTest(BaseApiTestCase):
             block_run_outputs_cache=None,
             cache_block_output_in_memory=False,
             callback_url=callback_url,
-            global_vars=global_vars,
+            global_vars=merge_dict(global_vars, dict(retry=dict(attempts=1))),
             input_from_output=input_from_output,
             logging_tags={
                 'block_type': BlockType.DBT,
@@ -167,9 +171,10 @@ class BlockExecutorTest(BaseApiTestCase):
         # )
         self.block_executor._execute_callback.assert_called_with(
             'on_success',
+            callback_kwargs=dict(retry=dict(attempts=1)),
             dynamic_block_index=None,
             dynamic_upstream_block_uuids=None,
-            global_vars=global_vars,
+            global_vars=ignore_keys(global_vars, ['retry']),
             logging_tags={
                 'block_type': BlockType.DBT,
                 'block_uuid': self.block_uuid,
@@ -243,7 +248,7 @@ class BlockExecutorTest(BaseApiTestCase):
 
         self.block_executor._execute_callback(
             callback='on_success',
-            global_vars={},
+            global_vars=dict(retry=dict(attempts=1)),
             logging_tags={},
             pipeline_run=None,
             dynamic_block_index=None,
@@ -255,7 +260,7 @@ class BlockExecutorTest(BaseApiTestCase):
             dynamic_block_index=None,
             dynamic_upstream_block_uuids=None,
             execution_partition=self.execution_partition,
-            global_vars={},
+            global_vars=dict(retry=dict(attempts=1)),
             logger=self.logger,
             logging_tags={},
             parent_block=self.block,
@@ -374,15 +379,21 @@ class BlockExecutorTest(BaseApiTestCase):
                 executor.execute(block_run_id=block_run.id, on_complete=__on_complete)
 
     def test_block_run_for_dynamic_child_block(self):
-        block1 = self.blocks[0]
+        pipeline, blocks = create_pipeline_with_blocks(
+            'dynamic_child_block',
+            self.repo_path,
+            return_blocks=True,
+        )
+
+        block1 = blocks[0]
         block1.configuration = dict(dynamic=True)
-        self.pipeline1.add_block(block1)
-        block2 = self.blocks[1]
+        pipeline.add_block(block1)
+        block2 = blocks[1]
 
         pipeline_run = PipelineRun.create(
             execution_date=datetime.utcnow(),
             pipeline_schedule_id=0,
-            pipeline_uuid=self.pipeline1.uuid,
+            pipeline_uuid=pipeline.uuid,
         )
         block_run = BlockRun.create(
             block_uuid=block2.uuid,
@@ -390,20 +401,18 @@ class BlockExecutorTest(BaseApiTestCase):
         )
 
         executor = BlockExecutor(
-            self.pipeline1,
+            pipeline,
             block_run.block_uuid,
+            block_run_id=block_run.id,
             execution_partition=pipeline_run.execution_partition,
         )
 
         self.assertTrue(isinstance(executor.block, DynamicChildBlockFactory))
-        self.assertTrue(executor.block.pipeline_run is None)
 
         with patch.object(executor.block, 'execute_sync') as mock_execute_sync:
             executor.execute(block_run_id=block_run.id, pipeline_run_id=pipeline_run.id)
             mock_execute_sync.assert_called_once()
 
-            self.assertEqual(executor.block.pipeline_run.id, pipeline_run.id)
-
         block_run_dynamic_child_spawn = BlockRun.create(
             block_uuid=f'{block2.uuid}:0',
             pipeline_run_id=pipeline_run.id,
@@ -412,46 +421,47 @@ class BlockExecutorTest(BaseApiTestCase):
             }))
         )
         executor = BlockExecutor(
-            self.pipeline1,
+            pipeline,
             block_run_dynamic_child_spawn.block_uuid,
             execution_partition=pipeline_run.execution_partition,
         )
 
         self.assertFalse(isinstance(executor.block, DynamicChildBlockFactory))
 
-        with patch.object(self.pipeline1, 'get_block', lambda _x: block2):
-            with patch.object(block2, 'execute_sync') as mock_execute_sync:
-                executor.execute(
+        # with patch.object(pipeline, 'get_block', lambda _x: block2):
+        with patch.object(block2, 'execute_sync') as mock_execute_sync:
+            executor.execute(
+                block_run_id=block_run_dynamic_child_spawn.id,
+                pipeline_run_id=pipeline_run.id,
+            )
+
+            mock_execute_sync.assert_called_once_with(
+                analyze_outputs=False,
+                block_run_outputs_cache=None,
+                dynamic_block_index=None,
+                dynamic_block_indexes=dict({
+                    block1.uuid: 0,
+                }),
+                dynamic_block_uuid=None,
+                dynamic_upstream_block_uuids=None,
+                execution_partition=pipeline_run.execution_partition,
+                global_vars=dict(retry=dict(attempts=1), logger=ANY),
+                input_from_output=None,
+                logger=ANY,
+                logging_tags=dict(
                     block_run_id=block_run_dynamic_child_spawn.id,
+                    block_type=block2.type,
+                    block_uuid=block_run_dynamic_child_spawn.block_uuid,
                     pipeline_run_id=pipeline_run.id,
-                )
-
-                mock_execute_sync.assert_called_once_with(
-                    analyze_outputs=False,
-                    block_run_outputs_cache=None,
-                    dynamic_block_index=None,
-                    dynamic_block_indexes=dict({
-                        block1.uuid: 0,
-                    }),
-                    dynamic_block_uuid=None,
-                    dynamic_upstream_block_uuids=None,
-                    execution_partition=pipeline_run.execution_partition,
-                    global_vars=None,
-                    input_from_output=None,
-                    logger=ANY,
-                    logging_tags=dict(
-                        block_run_id=block_run_dynamic_child_spawn.id,
-                        block_type=block2.type,
-                        block_uuid=block_run_dynamic_child_spawn.block_uuid,
-                        pipeline_run_id=pipeline_run.id,
-                        pipeline_uuid=self.pipeline1.uuid,
-                    ),
-                    run_all_blocks=True,
-                    runtime_arguments=None,
-                    store_variables=True,
-                    update_status=False,
-                    verify_output=True,
-                )
+                    pipeline_uuid=pipeline.uuid,
+                ),
+                metadata=None,
+                run_all_blocks=True,
+                runtime_arguments=None,
+                store_variables=True,
+                update_status=False,
+                verify_output=True,
+            )
 
     def test_block_run_for_dynamic_child_block_reduce_output(self):
         block1 = self.blocks[0]