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

mage_ai/shared/custom_logger.py

@@ -0,0 +1,213 @@
+import logging
+from enum import Enum
+
+import simplejson
+
+from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.shared.environments import is_deus_ex_machina
+from mage_ai.shared.parsers import encode_complex
+
+
+class Color(str, Enum):
+    BLUE = "\x1b[1;34m"
+    BOLD_RED = "\x1b[31;1m"
+    GREEN = "\x1b[1;32m"
+    GREY = "\x1b[38;20m"
+    LIGHT_BLUE = "\x1b[1;36m"
+    PINK = "\x1b[35m"
+    PURPLE = "\x1b[1;35m"
+    RED = "\x1b[31;20m"
+    RESET = "\x1b[0m"
+    YELLOW = "\x1b[33;21m"
+
+
+BLOCK_TYPE_COLOR_MAPPING = {
+    BlockType.CALLBACK: Color.GREEN,
+    BlockType.CONDITIONAL: Color.BLUE,
+    BlockType.CHART: Color.GREY,
+    BlockType.CUSTOM: Color.LIGHT_BLUE,
+    BlockType.DATA_EXPORTER: Color.YELLOW,
+    BlockType.DATA_LOADER: Color.BLUE,
+    BlockType.DBT: Color.RED,
+    BlockType.DYNAMIC_CHILD: Color.GREEN,
+    BlockType.EXTENSION: Color,
+    BlockType.GLOBAL_DATA_PRODUCT: Color.RED,
+    BlockType.HOOK: Color.BLUE,
+    BlockType.MARKDOWN: Color.GREY,
+    BlockType.SCRATCHPAD: Color.GREY,
+    BlockType.SENSOR: Color.PINK,
+    BlockType.TRANSFORMER: Color.PURPLE,
+}
+
+COLORS = [c for c in Color]
+
+
+def targets():
+    return [
+        # 'get_outputs',
+        # 'get_variables_by_block',
+        # 'fetch_input_variables',
+        # 'output_variables',
+        # 'uuid_for_output_variables',
+        # 'dynamic_block_values_and_metadata',
+    ]
+
+
+class ColorPrinter:
+    def __init__(self):
+        self.label = None
+
+    def debug(self, *args, **kwargs):
+        self.print(color=Color.PURPLE, *args, **kwargs)
+
+    def info(self, *args, **kwargs):
+        self.print(color=Color.BLUE, *args, **kwargs)
+
+    def warning(self, *args, **kwargs):
+        self.print(color=Color.GREEN, *args, **kwargs)
+
+    def error(self, *args, **kwargs):
+        self.print(color=Color.RED, *args, **kwargs)
+
+    def critical(self, *args, **kwargs):
+        self.print(color=Color.BOLD_RED, *args, **kwargs)
+
+    def print(self, *args, **kwargs):
+        if not is_deus_ex_machina():
+            return
+
+        __uuid = kwargs.pop('__uuid', None)
+        targs = targets()
+        if targs and __uuid not in targs:
+            return
+
+        color = kwargs.pop('color', None)
+        block = kwargs.pop('block', None)
+
+        more = None
+        if kwargs:
+            more = simplejson.dumps(
+                kwargs,
+                default=encode_complex,
+                ignore_nan=True,
+                indent=2,
+                sort_keys=True,
+            )
+            more = '\n'.join([f'\t\t{line}' for line in more.split('\n')])
+
+        output = ''.join(args)
+        if more:
+            output = f'{output}\n{more}'
+
+        label = __uuid or self.label
+        if label:
+            output = f'[{label}] {output}'
+
+        if not color:
+            color = COLORS[len(output) % len(COLORS)]
+
+        block_color = None
+        if block:
+            block_color = BLOCK_TYPE_COLOR_MAPPING.get(block.type) or color
+
+        print(f'{color}{"_" * 100}{color}')
+        print('\n')
+
+        if block:
+            print(f'{block_color}{"-" * 30}{block_color}')
+            print(f'{block_color}[Block] {block.uuid}{block_color}')
+            print(f'{block_color}\t\t Type    : {block.type}{block_color}')
+            print(f'{block_color}\t\t Language: {block.type}{block_color}')
+            print(f'{block_color}{"-" * 15}{block_color}')
+
+        print(f'{color}{output}{color}')
+
+        print('\n')
+        print(f'{color}{"_" * 100}{color}')
+
+
+class CustomFormatter(logging.Formatter):
+    # https://stackoverflow.com/a/56944256/1084284
+    format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
+
+    FORMATS = {
+        logging.CRITICAL: Color.BOLD_RED + format + Color.RESET,
+        logging.DEBUG: Color.PURPLE + format + Color.RESET,
+        logging.ERROR: Color.RED + format + Color.RESET,
+        logging.INFO: Color.BLUE + format + Color.RESET,
+        logging.NOTSET: Color.GREY + format + Color.RESET,
+        logging.WARNING: Color.GREEN + format + Color.RESET,
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+    def debug(self, p: int = 0, *args, **kwargs):
+        if is_deus_ex_machina():
+            if p:
+                self.print(color=Color.PURPLE, *args, **kwargs)
+            else:
+                super().debug(*args, **kwargs)
+
+    def info(self, p: int = 0, *args, **kwargs):
+        if is_deus_ex_machina():
+            if p:
+                self.print(color=Color.BLUE, *args, **kwargs)
+            else:
+                super().info(*args, **kwargs)
+
+    def warning(self, p: int = 0, *args, **kwargs):
+        if is_deus_ex_machina():
+            if p:
+                self.print(color=Color.YELLOW, *args, **kwargs)
+            else:
+                super().warning(*args, **kwargs)
+
+    def error(self, p: int = 0, *args, **kwargs):
+        if is_deus_ex_machina():
+            if p:
+                self.print(color=Color.RED, *args, **kwargs)
+            else:
+                super().error(*args, **kwargs)
+
+    def critical(self, p: int = 0, *args, **kwargs):
+        if is_deus_ex_machina():
+            if p:
+                self.print(color=Color.BOLD_RED, *args, **kwargs)
+            else:
+                super().critical(*args, **kwargs)
+
+    def print(self, color: Color, *args, **kwargs):
+        more = None
+        if kwargs:
+            more = simplejson.dumps(
+                kwargs,
+                default=encode_complex,
+                ignore_nan=True,
+                indent=2,
+                sort_keys=True,
+            )
+            more = '\n'.join([f'\t\t{line}' for line in more.split('\n')])
+
+        output = ''.join(args)
+        if more:
+            output = f'{output}\n{more}'
+
+        print(f'{color}{output}{color}')
+
+
+# create logger with 'spam_application'
+DX_LOGGER = logging.getLogger('deus_ex_machina')
+DX_LOGGER.setLevel(logging.DEBUG)
+
+# create console handler with a higher log level
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+
+ch.setFormatter(CustomFormatter())
+
+DX_LOGGER.addHandler(ch)
+
+DX_PRINTER = ColorPrinter()

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):