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

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

@@ -12,7 +12,6 @@ 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,
@@ -147,11 +146,12 @@ class ProjectPlatformAccessible:
 
         root_project_path, path, file_path_base = parts
 
-        return File(
-            filename=file_path_base,
-            dir_path=path,
-            repo_path=root_project_path,
-        )
+        if file_path_base and path and root_project_path:
+            return File(
+                filename=file_path_base,
+                dir_path=path,
+                repo_path=root_project_path,
+            )
 
     def hydrate_dbt_nodes(self, nodes_default: Dict, nodes_init: List[Dict]) -> Dict:
         """
@@ -173,7 +173,6 @@ class ProjectPlatformAccessible:
         #   /home/src/default_platform/default_repo/dbt/demo
         # node['original_file_path']
         #   models/example/model.sql
-        print('self.project_path', self.project_path)
         return {
             node['unique_id']: {
                 # file_path needs to be:
@@ -258,7 +257,7 @@ class ProjectPlatformAccessible:
             if paths:
                 return paths.get('full_path')
 
-        return get_repo_path(root_project=False)
+        return None
 
     def __file_source(self) -> str:
         return self.configuration.get('file_source') if self.configuration else None

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

@@ -9,20 +9,25 @@ 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.data_preparation.models.variables.constants import (
+    DATAFRAME_CSV_FILE,
+    VariableType,
+)
 from mage_ai.shared.parsers import encode_complex
 
 BLOCK_TYPE_TO_EXECUTION_TEMPLATE = {
-    BlockType.DATA_LOADER: 'data_loader.jinja',
-    BlockType.TRANSFORMER: 'transformer.jinja',
-    BlockType.DATA_EXPORTER: 'data_exporter.jinja',
+    BlockType.DATA_LOADER: "data_loader.jinja",
+    BlockType.TRANSFORMER: "transformer.jinja",
+    BlockType.DATA_EXPORTER: "data_exporter.jinja",
 }
 
 template_env = jinja2.Environment(
-    loader=jinja2.FileSystemLoader(os.path.join(
-        os.path.dirname(__file__),
-        'templates',
-    )),
+    loader=jinja2.FileSystemLoader(
+        os.path.join(
+            os.path.dirname(__file__),
+            "templates",
+        )
+    ),
     lstrip_blocks=True,
     trim_blocks=True,
 )
@@ -34,9 +39,12 @@ def execute_r_code(
     execution_partition: str = None,
     global_vars: Dict = None,
 ):
-    input_variable_objects = block.input_variable_objects(
-        execution_partition=execution_partition,
-    ) or []
+    input_variable_objects = (
+        block.input_variable_objects(
+            execution_partition=execution_partition,
+        )
+        or []
+    )
 
     # Render R script with user code
     execution_code = __render_r_script(
@@ -46,8 +54,8 @@ def execute_r_code(
         global_vars=global_vars,
         input_variable_objects=input_variable_objects,
     )
-    file_path = f'/tmp/{str(uuid.uuid4())}.r'
-    with open(file_path, 'w') as foutput:
+    file_path = f"/tmp/{str(uuid.uuid4())}.r"
+    with open(file_path, "w") as foutput:
         foutput.write(execution_code)
 
     # Convert input variable to csv format
@@ -60,14 +68,21 @@ def execute_r_code(
     output_variable_objects = block.output_variable_objects(
         execution_partition=execution_partition,
     )
-    output_variable_objects = [v for v in output_variable_objects
-                               if os.path.exists(os.path.join(
-                                    output_variable_objects[0].variable_path,
-                                    DATAFRAME_CSV_FILE,
-                                ))]
+    output_variable_objects = [
+        v
+        for v in output_variable_objects
+        if os.path.exists(
+            os.path.join(
+                output_variable_objects[0].variable_path,
+                DATAFRAME_CSV_FILE,
+            )
+        )
+    ]
 
     if len(output_variable_objects) > 0:
-        df = pd.read_csv(os.path.join(output_variable_objects[0].variable_path, DATAFRAME_CSV_FILE))
+        df = pd.read_csv(
+            os.path.join(output_variable_objects[0].variable_path, DATAFRAME_CSV_FILE)
+        )
     else:
         df = None
     return df
@@ -81,7 +96,7 @@ def __convert_inputs_to_csvs(input_variable_objects):
 
 def __render_global_vars(global_vars: Dict = None):
     if not global_vars:
-        return ''
+        return ""
 
     def format_value(val):
         if type(val) is int or type(val) is float:
@@ -96,9 +111,9 @@ def __render_global_vars(global_vars: Dict = None):
         else:
             return f"'{val}'"
 
-    var_list = [f'{k}={format_value(v)}' for k, v in global_vars.items()]
-    val_list_str = ', '.join(var_list)
-    return f'global_vars = c({val_list_str})'
+    var_list = [f"{k}={format_value(v)}" for k, v in global_vars.items()]
+    val_list_str = ", ".join(var_list)
+    return f"global_vars = c({val_list_str})"
 
 
 def __render_r_script(
@@ -112,12 +127,12 @@ def __render_r_script(
         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.',
+            f"Block execution for {block.type} with R language is not supported.",
         )
     template = template_env.get_template(BLOCK_TYPE_TO_EXECUTION_TEMPLATE[block.type])
 
     output_variable_object = block.variable_object(
-        'output_0',
+        "output_0",
         execution_partition=execution_partition,
     )
     os.makedirs(output_variable_object.variable_path, exist_ok=True)
@@ -125,23 +140,26 @@ def __render_r_script(
 
     global_vars_str = __render_global_vars(global_vars=global_vars)
 
-    return template.render(
-        code=code,
-        global_vars=global_vars_str,
-        input_paths=[os.path.join(v.variable_path, DATAFRAME_CSV_FILE)
-                     for v in input_variable_objects],
-        input_vars_str=', '.join([f'df_{i + 1}' for i in range(len(input_variable_objects))]),
-        output_path=output_path,
-    ) + '\n'
+    return (
+        template.render(
+            code=code,
+            global_vars=global_vars_str,
+            input_paths=[
+                os.path.join(v.variable_path, DATAFRAME_CSV_FILE)
+                for v in input_variable_objects
+            ],
+            input_vars_str=", ".join(
+                [f"df_{i + 1}" for i in range(len(input_variable_objects))]
+            ),
+            output_path=output_path,
+        )
+        + "\n"
+    )
 
 
 def __execute_r_code(file_path: str):
     subprocess.run(
-        [
-            'Rscript',
-            '--vanilla',
-            file_path
-        ],
+        ["Rscript", "--vanilla", file_path],
         check=True,
     )
 

mage_ai/data_preparation/models/block/settings/dynamic/constants.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+DEFAULT_STREAM_POLL_INTERVAL = 60
+
+
+class ModeType(str, Enum):
+    STREAM = 'stream'

mage_ai/data_preparation/models/block/settings/dynamic/mixins.py

@@ -0,0 +1,118 @@
+from typing import Optional
+
+from mage_ai.data_preparation.models.block.settings.dynamic.constants import ModeType
+from mage_ai.data_preparation.models.block.settings.dynamic.models import (
+    DynamicConfiguration,
+    ModeSettings,
+)
+from mage_ai.shared.array import find
+
+
+class DynamicMixin:
+    """
+    How the configuration looks in a pipeline’s metadata.yaml
+
+    - uuid: block_uuid
+      configuration:
+          dynamic:
+              parent: true
+    """
+
+    configuration = dict(dynamic=None)
+    downstream_blocks = []
+    replicated_block = None
+    upstream_blocks = []
+    uuid = None
+    uuid_replicated = None
+
+    def is_original_block(self, block_uuid: str) -> bool:
+        return self.uuid == block_uuid or (
+            self.replicated_block is not None and self.uuid_replicated == block_uuid
+        )
+
+    @property
+    def is_dynamic_v2(self) -> bool:
+        from mage_ai.settings.server import DYNAMIC_BLOCKS_V2
+
+        return DYNAMIC_BLOCKS_V2
+
+    @property
+    def is_dynamic_streaming(self) -> bool:
+        return self.is_dynamic_v2 and (
+            (self.is_dynamic_parent and self.is_dynamic_stream_mode_enabled)
+            or self.is_dynamic_child_streaming
+        )
+
+    @property
+    def is_dynamic_parent(self) -> bool:
+        return self.__dynamic_configuration().parent is not None
+
+    @property
+    def is_dynamic_child(self) -> bool:
+        return any(
+            upstream_block.should_dynamically_generate_block(self)
+            or upstream_block.is_dynamic_child
+            for upstream_block in self.upstream_blocks
+        )
+
+    @property
+    def is_dynamic_child_streaming(self) -> bool:
+        return self.is_dynamic_v2 and any(
+            (
+                upstream_block.should_dynamically_generate_block(self)
+                and upstream_block.is_dynamic_stream_mode_enabled
+            )
+            or upstream_block.is_dynamic_child_streaming
+            for upstream_block in self.upstream_blocks
+        )
+
+    @property
+    def should_reduce_output(self) -> bool:
+        return self.__dynamic_configuration().reduce_output is not None
+
+    @property
+    def is_dynamic_stream_mode_enabled(self) -> bool:
+        return self.settings_for_mode(ModeType.STREAM) is not None
+
+    def build_dynamic_uuid(self, index: int) -> str:
+        return ':'.join([self.uuid or '__missing_uuid__', str(index)])
+
+    def settings_for_mode(self, mode_type: ModeType) -> Optional[ModeSettings]:
+        modes = self.__dynamic_configuration().modes
+        if modes is not None:
+            return find(lambda ms: ms.type == mode_type, modes)
+
+    def should_reduce_output_from_upstream_block(self, block) -> bool:
+        reduce_output_upstream = self.__dynamic_configuration().reduce_output_upstream or []
+        return reduce_output_upstream is not None and block.uuid in reduce_output_upstream
+
+    def should_reduce_output_for_downstream_block(self, block) -> bool:
+        reduce_output = self.__dynamic_configuration().reduce_output
+
+        if reduce_output is True:
+            return True
+
+        if isinstance(reduce_output, list):
+            return block.uuid in reduce_output
+
+        return False
+
+    def should_dynamically_generate_block(self, block) -> bool:
+        parent = self.__dynamic_configuration().parent
+
+        if parent is True:
+            return True
+
+        if isinstance(parent, list):
+            return block.uuid in parent
+
+        return False
+
+    def __dynamic_configuration(self) -> DynamicConfiguration:
+        if self.configuration:
+            config = self.configuration.get('dynamic', {})
+            if config and isinstance(config, dict):
+                return DynamicConfiguration.load(**config)
+            elif isinstance(config, DynamicConfiguration):
+                return config
+        return DynamicConfiguration()

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

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.data_preparation.models.block.settings.dynamic.constants import ModeType
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class ModeSettings(BaseDataClass):
+    type: ModeType
+    poll_interval: Optional[int] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_enums('type', ModeType)
+
+
+@dataclass
+class DynamicConfiguration(BaseDataClass):
+    batch_settings: Optional[BatchSettings] = None
+    modes: Optional[List[ModeSettings]] = None
+    # Enable for all downstream or choose specific downstream blocks
+    parent: Optional[Union[List[str], bool]] = None
+    # Enable for all downstream or choose specific downstream blocks
+    reduce_output: Optional[Union[List[str], bool]] = None
+    # As a dynamic child, choose which upstream blocks to reduce their output
+    reduce_output_upstream: Optional[List[str]] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('batch_settings', BatchSettings)
+        self.serialize_attribute_classes('modes', ModeSettings)

mage_ai/data_preparation/models/block/settings/global_data_products/mixins.py

@@ -0,0 +1,20 @@
+from typing import Dict, Optional
+
+from mage_ai.data_preparation.models.block.settings.global_data_products.models import (
+    GlobalDataProductConfiguration,
+)
+
+
+class GlobalDataProductsMixin:
+    configuration = dict(global_data_products=None)
+
+    def global_data_products_configuration(
+        self,
+    ) -> Optional[Dict[str, GlobalDataProductConfiguration]]:
+        if self.configuration:
+            configs = self.configuration.get('global_data_products', {})
+            if configs and isinstance(configs, dict):
+                return {
+                    k: GlobalDataProductConfiguration.load(**v) if isinstance(v, dict) else v
+                    for k, v in configs.items()
+                }

mage_ai/data_preparation/models/block/settings/global_data_products/models.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional
+
+from mage_ai.shared.models import BaseDataClass
+
+
+class GlobalDataProductObjectType(str, Enum):
+    BLOCK = 'block'
+    PIPELINE = 'pipeline'
+
+
+@dataclass
+class OutdatedAfterSettings(BaseDataClass):
+    months: Optional[int] = None
+    seconds: Optional[int] = None
+    weeks: Optional[int] = None
+    years: Optional[int] = None
+
+
+@dataclass
+class OutdatedStartingAtSettings(BaseDataClass):
+    day_of_month: Optional[int] = None
+    day_of_week: Optional[int] = None
+    day_of_year: Optional[int] = None
+    hour_of_day: Optional[int] = None
+    minute_of_hour: Optional[int] = None
+    month_of_year: Optional[int] = None
+    second_of_minute: Optional[int] = None
+    week_of_month: Optional[int] = None
+    week_of_year: Optional[int] = None
+
+
+@dataclass
+class GlobalDataProductConfiguration(BaseDataClass):
+    uuid: str
+    object_type: Optional[GlobalDataProductObjectType] = None
+    object_uuid: Optional[str] = None
+    outdated_after: Optional[OutdatedAfterSettings] = None
+    outdated_starting_at: Optional[OutdatedStartingAtSettings] = None
+    project: Optional[str] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('outdated_after', OutdatedAfterSettings)
+        self.serialize_attribute_class('outdated_starting_at', OutdatedStartingAtSettings)
+        self.serialize_attribute_enum('object_type', GlobalDataProductObjectType)

mage_ai/data_preparation/models/block/settings/variables/mixins.py

@@ -0,0 +1,74 @@
+from typing import Dict, List, Optional
+
+from mage_ai.data.constants import InputDataType
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.data_preparation.models.block.settings.variables.models import (
+    ChunkKeyTypeUnion,
+    VariableConfiguration,
+    VariableSettings,
+)
+
+
+class VariablesMixin:
+    configuration = dict(variables=None)
+
+    @property
+    def variables_configuration(self) -> Optional[VariableConfiguration]:
+        if self.configuration:
+            configs = self.configuration.get('variables', {})
+            if configs and isinstance(configs, dict):
+                return VariableConfiguration.load(**configs)
+            elif isinstance(configs, VariableConfiguration):
+                return configs
+
+    @property
+    def variable_settings_upstream(self) -> Optional[Dict[str, VariableSettings]]:
+        if self.variables_configuration:
+            return self.variables_configuration.upstream
+
+    def upstream_batch_settings(self, block_uuid: str) -> BatchSettings:
+        return (
+            (self.variables_configuration or VariableConfiguration())
+            .upstream_settings(block_uuid)
+            .batch_settings
+        ) or BatchSettings()
+
+    def upstream_chunks(self, block_uuid: str) -> Optional[List[ChunkKeyTypeUnion]]:
+        return (
+            (self.variables_configuration or VariableConfiguration())
+            .upstream_settings(block_uuid)
+            .chunks
+        )
+
+    def input_data_types(self, block_uuid: str) -> List[InputDataType]:
+        return (
+            (self.variables_configuration or VariableConfiguration())
+            .upstream_settings(block_uuid)
+            .input_data_types
+        ) or [InputDataType.DEFAULT]
+
+    @property
+    def read_settings(self) -> VariableSettings:
+        settings = self.variables_configuration.read if self.variables_configuration else None
+        return settings or VariableSettings()
+
+    @property
+    def read_batch_settings(self) -> BatchSettings:
+        return self.read_settings.batch_settings or BatchSettings()
+
+    @property
+    def read_chunks(self) -> Optional[List[ChunkKeyTypeUnion]]:
+        return self.read_settings.chunks
+
+    @property
+    def write_settings(self) -> VariableSettings:
+        settings = self.variables_configuration.write if self.variables_configuration else None
+        return settings or VariableSettings()
+
+    @property
+    def write_batch_settings(self) -> BatchSettings:
+        return self.write_settings.batch_settings or BatchSettings()
+
+    @property
+    def write_chunks(self) -> Optional[List[ChunkKeyTypeUnion]]:
+        return self.write_settings.chunks

mage_ai/data_preparation/models/block/settings/variables/models.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+from mage_ai.data.constants import InputDataType
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.shared.models import BaseDataClass
+
+ChunkKeyTypeBase = Union[bool, float, int, str]
+ChunkKeyTypeDict = Dict[str, ChunkKeyTypeBase]
+ChunkKeyTypeUnion = Union[ChunkKeyTypeDict, ChunkKeyTypeBase]
+
+
+@dataclass
+class VariableSettings(BaseDataClass):
+    batch_settings: Optional[BatchSettings] = None
+    chunks: Optional[List[ChunkKeyTypeUnion]] = None
+    input_data_types: Optional[List[InputDataType]] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('batch_settings', BatchSettings)
+        self.serialize_attribute_enums('input_data_types', InputDataType)
+
+
+@dataclass
+class VariableConfiguration(BaseDataClass):
+    downstream: Optional[Dict[str, VariableSettings]] = None
+    read: Optional[VariableSettings] = None
+    upstream: Optional[Dict[str, VariableSettings]] = None
+    write: Optional[VariableSettings] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('read', VariableSettings)
+        self.serialize_attribute_class('write', VariableSettings)
+        self.downstream = {
+            k: VariableSettings.load(**v) if isinstance(v, dict) else v
+            for k, v in (self.downstream or {}).items()
+        }
+        self.upstream = {
+            k: VariableSettings.load(**v) if isinstance(v, dict) else v
+            for k, v in (self.upstream or {}).items()
+        }
+
+    def downstream_settings(self, block_uuid: str) -> VariableSettings:
+        settings = self.downstream.get(block_uuid) if self.downstream else None
+        return settings or VariableSettings()
+
+    def upstream_settings(self, block_uuid: str) -> VariableSettings:
+        settings = self.upstream.get(block_uuid) if self.upstream else None
+        return settings or VariableSettings()

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

@@ -57,7 +57,8 @@ class SparkBlock:
             )
 
     def compute_management_enabled(self) -> bool:
-        return Project(self.pipeline.repo_config if self.pipeline else None).is_feature_enabled(
+        return Project(
+            repo_config=self.pipeline.repo_config if self.pipeline else None).is_feature_enabled(
             FeatureUUID.COMPUTE_MANAGEMENT,
         )
 

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

@@ -21,6 +21,7 @@ from mage_ai.data_preparation.models.block.sql import (
     trino,
 )
 from mage_ai.data_preparation.models.block.sql.utils.shared import (
+    extract_full_table_name,
     has_create_or_insert_statement,
     has_drop_statement,
     has_update_statement,
@@ -31,7 +32,6 @@ from mage_ai.data_preparation.models.block.sql.utils.shared import (
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.io.base import QUERY_ROW_LIMIT, DataSource, ExportWritePolicy
 from mage_ai.io.config import ConfigFileLoader
-from mage_ai.settings.repo import get_repo_path
 
 PREVIEWABLE_BLOCK_TYPES = [
     BlockType.CUSTOM,
@@ -98,7 +98,7 @@ def execute_sql_code(
     disable_query_preprocessing = configuration.get('disable_query_preprocessing', False) or False
 
     if not config_file_loader:
-        config_path = path.join(get_repo_path(), 'io_config.yaml')
+        config_path = path.join(block.repo_path, 'io_config.yaml')
         config_profile = configuration.get('data_provider_profile')
         config_file_loader = ConfigFileLoader(config_path, config_profile)
 
@@ -662,6 +662,29 @@ def execute_raw_sql(
     should_query: bool = False,
     disable_query_preprocessing: bool = False,
 ) -> List[Any]:
+    """
+    Execute raw SQL queries with optional preprocessing and post-processing.
+
+    Args:
+        loader (Loader): The loader object responsible for executing queries.
+        block (Block): The SQL block object.
+        query_string (str): The SQL query string to execute.
+        configuration (Dict, optional): Additional configuration parameters for the query execution.
+            Defaults to None.
+        should_query (bool, optional): Flag indicating whether to perform query execution.
+            Defaults to False.
+        disable_query_preprocessing (bool, optional): Flag indicating whether to disable query
+            preprocessing. Defaults to False.
+
+    Returns:
+        List[Any]: A list containing the query results or an empty list if no results are expected.
+
+    Note:
+        This method preprocesses the query string, executes the queries, and performs
+        post-processing as needed. If `should_query` is False, it only preprocesses the query
+        without executing it.
+
+    """
     if configuration is None:
         configuration = {}
 
@@ -686,9 +709,11 @@ def execute_raw_sql(
             queries.append(query)
             fetch_query_at_indexes.append(True)
 
-    if should_query and (has_create_or_insert or has_update) and block.full_table_name:
-        queries.append(f'SELECT * FROM {block.full_table_name} LIMIT 1000')
-        fetch_query_at_indexes.append(block.full_table_name)
+    if should_query and (has_create_or_insert or has_update):
+        full_table_name = extract_full_table_name(query_string)
+        if full_table_name:
+            queries.append(f'SELECT * FROM {full_table_name} LIMIT 1000')
+            fetch_query_at_indexes.append(full_table_name)
 
     if should_query:
         results = loader.execute_queries(