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

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

@@ -1,6 +1,437 @@
-from typing import List
+import os
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Tuple
 
 from mage_ai.shared.array import find
+from mage_ai.shared.custom_logger import DX_PRINTER
+from mage_ai.shared.hash import ignore_keys_with_blank_values
+from mage_ai.shared.models import BaseDataClass
+
+
+class DynamicBlockFlag(str, Enum):
+    CLONE_OF_ORIGINAL = 'clone_of_original'
+    DYNAMIC = 'dynamic'
+    DYNAMIC_CHILD = 'dynamic_child'
+    ORIGINAL = 'original'
+    REDUCE_OUTPUT = 'reduce_output'
+    REPLICATED = 'replicated'
+    SPAWN_OF_DYNAMIC_CHILD = 'spawn_of_dynamic_child'
+
+
+def extract_dynamic_block_index(block_run_block_uuid: str) -> int:
+    if block_run_block_uuid:
+        parts = block_run_block_uuid.split(':')
+        if len(parts) >= 2:
+            return parts[-1]
+
+
+def is_dynamic_block(block) -> bool:
+    """
+    Checks if the given block is a dynamic block.
+
+    Args:
+        block: The block.
+
+    Returns:
+        bool: True if the block is a dynamic block, False otherwise.
+    """
+    return block.configuration and block.configuration.get('dynamic', False)
+
+
+def should_reduce_output(block) -> bool:
+    """
+    Checks if the given block should reduce its output.
+
+    Args:
+        block: The block.
+
+    Returns:
+        bool: True if the block should reduce its output, False otherwise.
+    """
+    if not block.configuration or not block.configuration.get('reduce_output', False):
+        return False
+    return True
+
+
+def has_dynamic_block_upstream_parent(block) -> bool:
+    return block.upstream_blocks and any([is_dynamic_block(b) for b in block.upstream_blocks])
+
+
+def is_dynamic_block_child(block) -> bool:
+    """
+    Checks if the given block is a dynamic block child.
+
+    Args:
+        block: The block.
+
+    Returns:
+        bool: True if the block is a dynamic block child, False otherwise.
+    """
+    dynamic_or_child = []
+
+    for upstream_block in block.upstream_blocks:
+        if is_dynamic_block(upstream_block) or is_dynamic_block_child(upstream_block):
+            dynamic_or_child.append(upstream_block)
+
+    if len(dynamic_or_child) == 0:
+        return False
+
+    dynamic_or_child_with_reduce = list(filter(lambda x: should_reduce_output(x), dynamic_or_child))
+
+    return len(dynamic_or_child) > len(dynamic_or_child_with_reduce)
+
+
+def is_replicated_block(block) -> bool:
+    return True if block and block.replicated_block else False
+
+
+def is_original_dynamic_child_block(
+    block,
+    block_run_block_uuid: int = None,
+    block_run_id: int = None,
+) -> bool:
+    # Check to see if its the original non-cloned version.
+    block_run = None
+
+    def __get_block_run(block_run_id=block_run_id):
+        from mage_ai.orchestration.db.models.schedules import BlockRun
+        return BlockRun.query.get(block_run_id)
+
+    if block:
+        if not block_run_block_uuid and block_run_id:
+            block_run = __get_block_run()
+            block_run_block_uuid = block_run.block_uuid
+
+        if block_run_block_uuid and block.uuid == block_run_block_uuid:
+            return True
+
+    if not block_run and not block_run_id:
+        return False
+
+    if not block_run:
+        block_run = __get_block_run()
+
+    if block and block.uuid == block_run.block_uuid:
+        return True
+
+    wrapper = DynamicBlockWrapperBase()
+    wrapper.hydrate(block=block, block_run=block_run)
+
+    return wrapper.is_dynamic_child() and wrapper.is_original(include_clone=True)
+
+
+def uuid_for_output_variables(
+    block,
+    block_uuid: str = None,
+    dynamic_block_index: int = None,
+    dynamic_block_uuid: str = None,
+    join_character: str = None,
+    **kwargs,
+) -> Tuple[str, bool]:
+    changed = False
+
+    block_uuid_0 = block_uuid
+    dynamic_block_index_0 = dynamic_block_index
+
+    if block_uuid is None:
+        block_uuid = block.uuid
+
+    if dynamic_block_uuid:
+        block_uuid = dynamic_block_uuid
+        dynamic_block_index = None
+
+    if dynamic_block_index is not None or is_dynamic_block_child(block):
+        parts = block_uuid.split(':')
+
+        if len(parts) >= 2:
+            block_uuid = parts[0]
+
+            if dynamic_block_index is None:
+                dynamic_block_index = parts[-1]
+
+        if dynamic_block_index is not None:
+            # We only need the base name and the final index to create the folder structure:
+            # e.g. block_uuid/[dynamic_block_index]
+            # e.g. block_uuid/0/output_0/data.json
+            # [dynamic_block_index] if used for each dynamic child so that it has a folder
+            # to store its output.
+            # dynamic_block_index = dynamic_block_uuid.split(':')[-1]
+            # block_uuid = os.path.join(block_uuid, str(dynamic_block_index))
+            arr = [str(block_uuid), str(dynamic_block_index)]
+            if join_character:
+                block_uuid = join_character.join(arr)
+            else:
+                block_uuid = os.path.join(*arr)
+            changed = True
+
+    DX_PRINTER.debug(
+        block=block,
+        block_uuid=block_uuid,
+        block_uuid_0=block_uuid_0,
+        changed=changed,
+        dynamic_block_index=dynamic_block_index,
+        dynamic_block_index_0=dynamic_block_index_0,
+        dynamic_block_uuid=dynamic_block_uuid,
+        __uuid='uuid_for_output_variables',
+    )
+
+    return (block_uuid, changed)
+
+
+@dataclass
+class DynamicBlockWrapperBase(BaseDataClass):
+    block: Any = None
+    block_run_block_uuid: str = None
+    block_uuid: str = None
+    children: List[BaseDataClass] = field(default_factory=lambda: [])
+    clones: List[BaseDataClass] = field(default_factory=lambda: [])
+    dynamic_block_index: int = None
+    dynamic_block_indexes: Dict = None
+    factory: Any = None
+    flags: List[DynamicBlockFlag] = field(default_factory=lambda: [])
+    index: int = None
+    metadata_instructions: BaseDataClass = field(default_factory=lambda: {})
+    siblings: List[BaseDataClass] = field(default_factory=lambda: [])
+    spawns: List[BaseDataClass] = field(default_factory=lambda: [])
+    upstream_dynamic_blocks: List[BaseDataClass] = field(default_factory=lambda: [])
+    upstream_dynamic_child_blocks: List[BaseDataClass] = field(default_factory=lambda: [])
+    uuid: str = None
+
+
+@dataclass
+class DynamicBlockWrapper(BaseDataClass):
+    block: Any = None
+    block_run_block_uuid: str = None
+    block_uuid: str = None
+    # If the current block is a dynamic block, then the children are the dynamic child blocks
+    # that are created from this block’s output.
+    children: List[DynamicBlockWrapperBase] = field(default_factory=lambda: [])
+    # If the current block is the original, then clones are the other blocks that the original
+    # created to handle the responsibility of controlling what block runs are created
+    # from an upstream dynamic block.
+    clones: List[DynamicBlockWrapperBase] = field(default_factory=lambda: [])
+    dynamic_block_index: int = None
+    dynamic_block_indexes: Dict = None
+    factory: Any = None
+    flags: List[DynamicBlockFlag] = field(default_factory=lambda: [])
+    metadata_instructions: BaseDataClass = None
+    # Other blocks at the same level; other cloned blocks, other spawns, etc.
+    siblings: List[DynamicBlockWrapperBase] = field(default_factory=lambda: [])
+    # If the current block is an original dynamic child block or cloned dynamic child block,
+    # then spawns are the blocks that are created based on the original or cloned upstream
+    # dynamic block’s (e.g. a direct parent block) output.
+    spawns: List[DynamicBlockWrapperBase] = field(default_factory=lambda: [])
+    # Upstream blocks that are dynamic blocks.
+    upstream_dynamic_blocks: List[DynamicBlockWrapperBase] = field(default_factory=lambda: [])
+    # Upstream blocks that are dynamic child blocks.
+    upstream_dynamic_child_blocks: List[DynamicBlockWrapperBase] = field(default_factory=lambda: [])
+    # Unique identifier used as a suffix in the block run.
+    uuid: str = None
+
+    def __post_init__(self):
+        self.factory._wrapper = self
+
+        block_run = self.factory.block_run()
+
+        self.block = self.factory.block
+        self.block_uuid = self.block.uuid if self.block else None
+        self.block_run_block_uuid = block_run.block_uuid if block_run else None
+        self.uuid = self.block_run_block_uuid or self.block_uuid
+
+        if block_run:
+            config = block_run.metrics or {}
+            self.dynamic_block_index = config.get('dynamic_block_index')
+            self.dynamic_block_indexes = config.get('dynamic_block_indexes')
+
+            metadata = config.get('metadata') or {}
+
+            self.flags = [DynamicBlockFlag(v) for v in metadata.get('flags') or []]
+            if metadata.get('clone_original', False):
+                self.flags.append(DynamicBlockFlag.CLONE_OF_ORIGINAL)
+            self.flags = list(set(self.flags))
+
+            for key in [
+                'uuid',
+            ]:
+                value = metadata.get(key) or None
+                if value:
+                    setattr(self, key, value)
+
+            for key in [
+                'children',
+                'clones',
+                'siblings',
+                'spawns',
+                'upstream_dynamic_blocks',
+                'upstream_dynamic_child_blocks',
+            ]:
+                values = metadata.get(key) or None
+                if values:
+                    setattr(self, key, [self.load(**m) for m in values])
+
+    def is_original(self, include_clone: bool = False) -> bool:
+        if DynamicBlockFlag.ORIGINAL in (self.flags or []) or (
+            include_clone and self.is_clone_of_original()
+        ):
+            return True
+
+        if not self.block:
+            return False
+
+        return self.block.uuid == self.block_run_block_uuid or \
+            (self.is_replicated() and self.block.uuid_replicated == self.block_run_block_uuid)
+
+    def is_replicated(self) -> bool:
+        return DynamicBlockFlag.REPLICATED in (self.flags or []) or is_replicated_block(self.block)
+
+    def is_clone_of_original(self) -> bool:
+        return DynamicBlockFlag.CLONE_OF_ORIGINAL in (self.flags or []) and not self.is_original()
+
+    def is_dynamic(self) -> bool:
+        if DynamicBlockFlag.DYNAMIC in (self.flags or []):
+            return True
+        return self.block and is_dynamic_block(self.block)
+
+    def is_dynamic_child(self) -> bool:
+        if DynamicBlockFlag.DYNAMIC_CHILD in (self.flags or []):
+            return True
+        return self.block and is_dynamic_block_child(self.block)
+
+    def is_dynamic_squared(self) -> bool:
+        return self.is_dynamic() and self.is_dynamic_child()
+
+    def is_spawn(self) -> bool:
+        if DynamicBlockFlag.SPAWN_OF_DYNAMIC_CHILD in (self.flags or []):
+            return True
+
+        return (self.block or self.block_uuid) and \
+            self.block_run_block_uuid and \
+            ((self.block and self.block.uuid != self.block_run_block_uuid) or
+                (self.block_uuid and self.block_uuid != self.block_run_block_uuid)) and \
+            not self.is_original(include_clone=True)
+
+    def should_reduce_output(self) -> bool:
+        if DynamicBlockFlag.REDUCE_OUTPUT in (self.flags or []):
+            return True
+        return self.block and should_reduce_output(self.block)
+
+    def get_dynamic_block_index_from_parent(
+        self,
+        block_run_block_uuid: str,
+    ) -> int:
+        if not self.dynamic_block_indexes:
+            block_run = self.factory.block_run()
+            if block_run and block_run.metrics:
+                self.dynamic_block_indexes = (block_run.metrics or {}).get(
+                    'dynamic_block_indexes',
+                )
+
+        if self.dynamic_block_indexes:
+            return (self.dynamic_block_indexes or {}).get(block_run_block_uuid)
+
+    def get_parent_block_uuid_for_output_variables(
+        self,
+        block,
+        block_run_block_uuid: str,
+    ) -> str:
+        dynamic_block_index_from_parent = self.get_dynamic_block_index_from_parent(
+            block_run_block_uuid,
+        )
+
+        return uuid_for_output_variables(
+            block,
+            block_uuid=block_run_block_uuid,
+            dynamic_block_index=dynamic_block_index_from_parent,
+            join_character=':',
+        )[0]
+
+    def get_dynamic_block_index(self) -> int:
+        if self.dynamic_block_index is not None:
+            return self.dynamic_block_index
+
+        block_run = self.factory.block_run()
+        if block_run and block_run.metrics:
+            self.dynamic_block_index = (block_run.metrics or {}).get(
+                'dynamic_block_index',
+            )
+
+        return self.dynamic_block_index
+
+    def to_dict_base(self, **kwargs) -> dict:
+        data = dict(
+            dynamic_block_index=self.get_dynamic_block_index(),
+        )
+
+        flags = []
+        if self.is_original():
+            flags.append(DynamicBlockFlag.ORIGINAL)
+        if self.is_clone_of_original():
+            flags.append(DynamicBlockFlag.CLONE_OF_ORIGINAL)
+        if self.is_dynamic():
+            flags.append(DynamicBlockFlag.DYNAMIC)
+        if self.is_dynamic_child():
+            flags.append(DynamicBlockFlag.DYNAMIC_CHILD)
+        if self.is_spawn():
+            flags.append(DynamicBlockFlag.SPAWN_OF_DYNAMIC_CHILD)
+        if self.should_reduce_output():
+            flags.append(DynamicBlockFlag.REDUCE_OUTPUT)
+
+        flags = list(set(flags + (self.flags or [])))
+        if len(flags) >= 1:
+            data['flags'] = [v.value for v in flags if isinstance(v, DynamicBlockFlag)]
+
+        for key in [
+            'block_run_block_uuid',
+            'block_uuid',
+            'uuid',
+        ]:
+            if getattr(self, key) is not None:
+                data[key] = getattr(self, key)
+
+        return ignore_keys_with_blank_values(data)
+
+    def to_dict(
+        self,
+        include_all: bool = False,
+        use_metadata_instructions: bool = False,
+        **kwargs,
+    ) -> dict:
+        if use_metadata_instructions:
+            return ignore_keys_with_blank_values(dict(
+                clone_original=self.metadata_instructions.clone_original,
+                original=self.metadata_instructions.original.to_dict_base(),
+                upstream=self.metadata_instructions.upstream.to_dict_base(),
+            ))
+
+        data = self.to_dict_base(**kwargs)
+        if include_all:
+            for key in [
+                'children',
+                'clones',
+                'siblings',
+                'spawns',
+                'upstream_dynamic_blocks',
+                'upstream_dynamic_child_blocks',
+            ]:
+                values = getattr(self, key) or None
+                if values:
+                    data[key] = [v.to_dict_base(**kwargs) for v in values]
+
+        return ignore_keys_with_blank_values(data)
+
+
+@dataclass
+class MetadataInstructions(BaseDataClass):
+    clone_original: bool = False
+    original: DynamicBlockWrapper = None
+    parent: DynamicBlockWrapper = None
+    upstream: DynamicBlockWrapper = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('original', DynamicBlockWrapper)
+        self.serialize_attribute_class('parent', DynamicBlockWrapper)
+        self.serialize_attribute_class('upstream', DynamicBlockWrapper)
 
 
 def dynamically_created_child_block_runs(pipeline, block, block_runs: List):

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

@@ -1,20 +1,25 @@
 import os
 from typing import Dict, List
 
+from mage_ai.cache.dbt.utils import get_project_path_from_file_path
 from mage_ai.data_preparation.models.block.platform.utils import (
     from_another_project,
     get_selected_directory_from_file_path,
 )
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.file import File
-from mage_ai.settings.platform import project_platform_activated
+from mage_ai.settings.platform import (
+    get_repo_paths_for_file_path,
+    project_platform_activated,
+)
+from mage_ai.settings.repo import get_repo_path
 from mage_ai.settings.utils import base_repo_path
 from mage_ai.shared.path_fixer import (
+    add_absolute_path,
     add_root_repo_path_to_relative_path,
     get_path_parts,
     remove_base_repo_directory_name,
     remove_base_repo_name,
-    remove_base_repo_path,
 )
 from mage_ai.shared.strings import remove_extension_from_filename
 from mage_ai.shared.utils import clean_name
@@ -36,51 +41,48 @@ class ProjectPlatformAccessible:
         if not config:
             return config
 
-        if self.project_platform_activated:
-            if config.get('file_source'):
-                file_source = config.get('file_source') or {}
-                if file_source:
-                    if file_source.get('path'):
-                        # default_platform/tons_of_dbt_projects/diff_name/models/example/
-                        # my_first_dbt_model.sql
-                        path = file_source.get('path')
-                        path = str(remove_base_repo_name(path)) if path else path
-                        file_source['path'] = path
-                        config['file_source'] = file_source
-
-                        if path and \
-                                BlockType.DBT == self.type and \
-                                BlockLanguage.YAML != self.language:
-
-                            # /home/src/default_repo/default_platform/
-                            # tons_of_dbt_projects/diff_name
-                            project_path = get_selected_directory_from_file_path(
-                                file_path=path,
-                                selector=lambda fn: (
-                                    str(fn).endswith('dbt_project.yml') or
-                                    str(fn).endswith('dbt_project.yaml')
-                                ),
+        if config.get('file_source'):
+            file_source = config.get('file_source') or {}
+            if file_source:
+                if file_source.get('path'):
+                    # default_platform/tons_of_dbt_projects/diff_name/models/example/
+                    # my_first_dbt_model.sql
+                    path = file_source.get('path')
+                    path = add_absolute_path(path, add_base_repo_path=False)
+                    file_source['path'] = path
+                    config['file_source'] = file_source
+
+                    if path and \
+                            BlockType.DBT == self.type and \
+                            BlockLanguage.YAML != self.language:
+
+                        # /home/src/default_repo/default_platform/
+                        # tons_of_dbt_projects/diff_name
+                        project_path = get_selected_directory_from_file_path(
+                            file_path=path,
+                            selector=lambda fn: (
+                                str(fn).endswith('dbt_project.yml') or
+                                str(fn).endswith('dbt_project.yaml')
+                            ),
+                        )
+                        # tons_of_dbt_projects/diff_name
+                        if project_path:
+                            file_source['project_path'] = add_absolute_path(
+                                project_path,
+                                add_base_repo_path=False
                             )
-                            # tons_of_dbt_projects/diff_name
-                            if project_path:
-                                file_source['project_path'] = remove_base_repo_path(
-                                    project_path,
-                                )
-
-            if config.get('file_path'):
-                file_path = config.get('file_path')
-                config['file_path'] = str(remove_base_repo_name(
-                    file_path,
-                )) if file_path else file_path
-        else:
-            if config.get('file_source'):
-                config.pop('file_source', None)
+        elif config.get('file_path'):
+            file_path = config.get('file_path')
+            config['file_path'] = str(add_absolute_path(
+                file_path,
+                add_base_repo_path=False
+            )) if file_path else file_path
 
         return config
 
     def is_from_another_project(self) -> bool:
         return self.project_platform_activated and from_another_project(
-            self.__file_source_path(),
+            self.file_source_path(),
             other_file_path=self.pipeline.dir_path if self.pipeline else None,
         )
 
@@ -88,7 +90,7 @@ class ProjectPlatformAccessible:
         if not self.project_platform_activated:
             return
 
-        return self.__file_source_path()
+        return self.file_source_path()
 
     def get_project_path_from_source(self) -> str:
         if not self.project_platform_activated:
@@ -98,6 +100,10 @@ class ProjectPlatformAccessible:
         project_path_relative = self.__file_source_project()
         if project_path_relative:
             return os.path.join(base_repo_path(), project_path_relative)
+        else:
+            file_path = self.file_source_path() or self.configuration.get('file_path')
+            if file_path:
+                return get_project_path_from_file_path(file_path, walk_up_parents=True)
 
     def get_project_path_from_project_name(self, project_name: str) -> str:
         if not self.project_platform_activated or not project_name:
@@ -116,7 +122,7 @@ class ProjectPlatformAccessible:
             )
 
             # A project profile hasn’t been selected yet for this block (it’s done through the UI).
-            root_project_path, path, file_path_base = get_path_parts(self.__file_source_path())
+            root_project_path, path, file_path_base = get_path_parts(self.file_source_path())
             project_name = os.path.join(path, DBT_DIRECTORY_NAME)
 
         return add_root_repo_path_to_relative_path(project_name)
@@ -236,10 +242,28 @@ class ProjectPlatformAccessible:
             pipeline=pipeline,
         )
 
+    def get_repo_path_from_configuration(self) -> str:
+        # Example:
+        # default_repo/dbt/demo/models/example/model.sql
+        # main_app/platform/dbt/demo/models/model.sql
+
+        file_path = self.get_file_path_from_source()
+        if file_path:
+            file_path = add_root_repo_path_to_relative_path(file_path)
+        elif self.configuration.get('file_path'):
+            file_path = self.configuration.get('file_path')
+
+        if file_path:
+            paths = get_repo_paths_for_file_path(file_path)
+            if paths:
+                return paths.get('full_path')
+
+        return get_repo_path(root_project=False)
+
     def __file_source(self) -> str:
         return self.configuration.get('file_source') if self.configuration else None
 
-    def __file_source_path(self) -> str:
+    def file_source_path(self) -> str:
         file_source = self.__file_source()
         return file_source.get('path') if file_source else None
 

mage_ai/data_preparation/models/block/r/__init__.py

@@ -1,15 +1,16 @@
-from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.constants import BlockType
-from mage_ai.data_preparation.models.variable import (
-    DATAFRAME_CSV_FILE,
-    VariableType,
-)
-from typing import Dict, List
-import jinja2
 import os
 import subprocess
 import uuid
+from typing import Dict, List
+
+import jinja2
 import pandas as pd
+import simplejson
+
+from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.data_preparation.models.variable import DATAFRAME_CSV_FILE, VariableType
+from mage_ai.shared.parsers import encode_complex
 
 BLOCK_TYPE_TO_EXECUTION_TEMPLATE = {
     BlockType.DATA_LOADER: 'data_loader.jinja',
@@ -85,6 +86,13 @@ def __render_global_vars(global_vars: Dict = None):
     def format_value(val):
         if type(val) is int or type(val) is float:
             return val
+        elif type(val) is dict:
+            json_val = simplejson.dumps(
+                val,
+                default=encode_complex,
+                ignore_nan=True,
+            )
+            return f"'{json_val}'"
         else:
             return f"'{val}'"
 
@@ -98,8 +106,10 @@ def __render_r_script(
     code: str,
     execution_partition: str = None,
     global_vars: Dict = None,
-    input_variable_objects: List = [],
+    input_variable_objects: List = None,
 ):
+    if input_variable_objects is None:
+        input_variable_objects = []
     if block.type not in BLOCK_TYPE_TO_EXECUTION_TEMPLATE:
         raise Exception(
             f'Block execution for {block.type} with R language is not supported.',