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

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(

mage_ai/data_preparation/models/block/sql/utils/shared.py

@@ -1,6 +1,6 @@
 import re
 from os import path
-from typing import Callable, Dict, List, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from jinja2 import StrictUndefined, Template
 from pandas import DataFrame
@@ -482,7 +482,7 @@ def remove_comments(text: str) -> str:
     return '\n'.join(line for line in lines if not line.startswith('--'))
 
 
-def extract_create_statement_table_name(text: str) -> str:
+def extract_create_statement_table_name(text: str) -> Optional[str]:
     create_table_pattern = r'create table(?: if not exists)*'
 
     statement_partial, _ = extract_and_replace_text_between_strings(
@@ -507,7 +507,7 @@ def extract_create_statement_table_name(text: str) -> str:
 
 def extract_insert_statement_table_names(text: str) -> List[str]:
     matches = re.findall(
-        r'insert(?: overwrite)*(?: into)*[\s]+([\w.]+)',
+        r'insert(?:\s+ignore)?(?:\s+overwrite)?(?:\s+into)?[\s]+([\w.]+)',
         remove_comments(text),
         re.IGNORECASE,
     )
@@ -532,6 +532,24 @@ def extract_update_statement_table_names(text: str) -> List[str]:
     return matches
 
 
+def extract_full_table_name(text: str) -> Optional[str]:
+    if not text:
+        return None
+
+    table_name = extract_create_statement_table_name(text)
+    if table_name:
+        return table_name
+
+    matches = extract_insert_statement_table_names(text)
+    if len(matches) == 0:
+        matches = extract_update_statement_table_names(text)
+
+    if len(matches) == 0:
+        return None
+
+    return matches[len(matches) - 1]
+
+
 def has_create_or_insert_statement(text: str) -> bool:
     table_name = extract_create_statement_table_name(text)
     if table_name:

mage_ai/data_preparation/models/block/utils.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import pandas as pd
 import simplejson
@@ -29,6 +29,7 @@ from mage_ai.data_preparation.models.block.dynamic.utils import (
 from mage_ai.data_preparation.models.block.dynamic.utils import (
     should_reduce_output as should_reduce_output_original,
 )
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.constants import (
     DATAFRAME_ANALYSIS_MAX_COLUMNS,
     BlockType,
@@ -125,10 +126,12 @@ def get_all_ancestors(block) -> List:
         List: A list of all ancestors of the block.
     """
     arr = get_leaf_nodes(block, 'upstream_blocks', include_all_nodes=True)
-    return list(filter(
-        lambda x: x.uuid != block.uuid,
-        arr,
-    ))
+    return list(
+        filter(
+            lambda x: x.uuid != block.uuid,
+            arr,
+        )
+    )
 
 
 def get_all_descendants(block) -> List:
@@ -142,10 +145,12 @@ def get_all_descendants(block) -> List:
         List: A list of all descendants of the block.
     """
     arr = get_leaf_nodes(block, 'downstream_blocks', include_all_nodes=True)
-    return list(filter(
-        lambda x: x.uuid != block.uuid,
-        arr,
-    ))
+    return list(
+        filter(
+            lambda x: x.uuid != block.uuid,
+            arr,
+        )
+    )
 
 
 def get_leaf_nodes(
@@ -195,6 +200,7 @@ def output_variables(
     include_df: bool = True,
     input_args: List[Any] = None,
     data_integration_settings_mapping: Dict = None,
+    max_results: Optional[int] = None,
 ) -> List[str]:
     """
     Retrieve the output variables generated by a specified block within a pipeline.
@@ -221,8 +227,9 @@ def output_variables(
 
     di_settings = None
     if data_integration_settings_mapping:
-        di_settings = data_integration_settings_mapping.get(block_uuid) or \
-            data_integration_settings_mapping.get(block.uuid)
+        di_settings = data_integration_settings_mapping.get(
+            block_uuid
+        ) or data_integration_settings_mapping.get(block.uuid)
 
     if block and not di_settings:
         di_settings = block.get_data_integration_settings(
@@ -238,17 +245,18 @@ def output_variables(
     if should_reduce_output(block):
         all_variables = all_variable_uuids(
             block,
+            max_results=max_results,
             partition=execution_partition,
         )
     else:
         all_variables = block.get_variables_by_block(
             block_uuid=block_uuid,
             dynamic_block_index=dynamic_block_index,
+            max_results=max_results,
             partition=execution_partition,
         )
 
-    output_variables = [v for v in all_variables
-                        if is_output_variable(v, include_df=include_df)]
+    output_variables = [v for v in all_variables if is_output_variable(v, include_df=include_df)]
 
     DX_PRINTER.error(
         block=block,
@@ -393,6 +401,7 @@ def fetch_input_variables(
     dynamic_block_flags: List[DynamicBlockFlag] = None,
     metadata: Dict = None,
     upstream_block_uuids_override: List[str] = None,
+    current_block=None,
 ) -> Tuple[List, List, List]:
     """
     Fetches the input variables for a block.
@@ -418,8 +427,9 @@ def fetch_input_variables(
         upstream_block_uuids = upstream_block_uuids_override
 
     kwargs_vars = []
-
     input_vars = []
+    input_variables_by_uuid = None
+
     if input_args is not None:
         input_vars = input_args
         if upstream_block_uuids and not upstream_block_uuids_override:
@@ -432,16 +442,21 @@ def fetch_input_variables(
         disable_dynamic_index_for_output_variables = False
         # Fetch the data normally, then use the dynamic block index to select.
         block_is_dynamic_block = DynamicBlockFlag.DYNAMIC in (dynamic_block_flags or [])
-        block_is_dynamic_block_child = DynamicBlockFlag.DYNAMIC_CHILD in (dynamic_block_flags or [])
+        block_is_dynamic_block_child = DynamicBlockFlag.DYNAMIC_CHILD in (
+            dynamic_block_flags or []
+        )
         if block_is_dynamic_block and block_is_dynamic_block_child:
             disable_dynamic_index_for_output_variables = True
 
         if dynamic_block_indexes and (
-            not block_is_dynamic_block or
-            not block_is_dynamic_block_child or
-            len(dynamic_block_indexes) >= 2
+            not block_is_dynamic_block
+            or not block_is_dynamic_block_child
+            or len(dynamic_block_indexes) >= 2
         ):
-            for upstream_block_uuid, dynamic_block_index_folder in dynamic_block_indexes.items():
+            for (
+                upstream_block_uuid,
+                dynamic_block_index_folder,
+            ) in dynamic_block_indexes.items():
                 block = pipeline.get_block(upstream_block_uuid)
                 if block:
                     upstream_is_dynamic_block = is_dynamic_block(block)
@@ -453,11 +468,12 @@ def fetch_input_variables(
                     - Upstream is a dynamic child block
                     - Current block is a virtual clone of the original one
                     """
-                    if upstream_is_dynamic_block and \
-                            upstream_is_dynamic_block_child and \
-                            metadata and \
-                            metadata.get('clone_original'):
-
+                    if (
+                        upstream_is_dynamic_block
+                        and upstream_is_dynamic_block_child
+                        and metadata
+                        and metadata.get('clone_original')
+                    ):
                         dynamic_block_index_mapping[block.uuid] = extract_dynamic_block_index(
                             upstream_block_uuid,
                         )
@@ -491,8 +507,9 @@ def fetch_input_variables(
             input_args=input_args,
             data_integration_settings_mapping=data_integration_settings_mapping,
             **(
-                {} if disable_dynamic_index_for_output_variables else
-                dynamic_block_index_values_for_output_variables
+                {}
+                if disable_dynamic_index_for_output_variables
+                else dynamic_block_index_values_for_output_variables
             ),
         )
         # Block UUIDs
@@ -510,33 +527,55 @@ def fetch_input_variables(
             if BlockType.GLOBAL_DATA_PRODUCT == upstream_block.type:
                 global_data_product = upstream_block.get_global_data_product()
                 input_vars[idx] = global_data_product.get_outputs()
+
+                mds = {}
+                variable_uuids = upstream_block.output_variables(
+                    execution_partition=execution_partition,
+                )
+                for variable_uuid in variable_uuids:
+                    md = pipeline.get_block_variable(
+                        upstream_block_uuid,
+                        variable_uuid,
+                        partition=execution_partition,
+                    )
+                    if isinstance(md, dict):
+                        mds.update(md)
+
+                kwargs_vars.append(mds)
+
                 continue
 
             # Block output variables for upstream_block_uuid
             variables = input_variables_by_uuid[upstream_block_uuid]
 
-            dynamic_block_index_for_output_variable = \
+            dynamic_block_index_for_output_variable = (
                 dynamic_block_index_values_for_output_variables.get('dynamic_block_index')
+            )
 
-            if not disable_dynamic_index_for_output_variables and \
-                    dynamic_block_index_mapping and \
-                    upstream_block_uuid in dynamic_block_index_mapping:
-
-                dynamic_block_index_for_output_variable = \
-                    dynamic_block_index_mapping[upstream_block_uuid]
+            if (
+                not disable_dynamic_index_for_output_variables
+                and dynamic_block_index_mapping
+                and upstream_block_uuid in dynamic_block_index_mapping
+            ):
+                dynamic_block_index_for_output_variable = dynamic_block_index_mapping[
+                    upstream_block_uuid
+                ]
 
             # Fetch variable values
             if should_reduce:
-                variable_values = [reduce_output_from_block(
-                    upstream_block,
-                    variable_uuid,
-                    from_notebook=from_notebook,
-                    global_vars=global_vars,
-                    input_args=input_args,
-                    partition=execution_partition,
-                    raise_exception=True,
-                    spark=spark,
-                ) for variable_uuid in variables]
+                variable_values = [
+                    reduce_output_from_block(
+                        upstream_block,
+                        variable_uuid,
+                        from_notebook=from_notebook,
+                        global_vars=global_vars,
+                        input_args=input_args,
+                        partition=execution_partition,
+                        raise_exception=True,
+                        spark=spark,
+                    )
+                    for variable_uuid in variables
+                ]
             else:
                 # Getting input variables from cache the cache is not empty
                 if block_run_outputs_cache:
@@ -553,6 +592,17 @@ def fetch_input_variables(
                             raise_exception=True,
                             spark=spark,
                             dynamic_block_index=dynamic_block_index_for_output_variable,
+                            read_chunks=current_block.upstream_chunks(upstream_block_uuid)
+                            if current_block
+                            else None,
+                            read_batch_settings=current_block.upstream_batch_settings(
+                                upstream_block_uuid
+                            )
+                            if current_block
+                            else None,
+                            input_data_types=current_block.input_data_types(upstream_block_uuid)
+                            if current_block
+                            else None,
                         )
                         for var in variables
                     ]
@@ -566,10 +616,11 @@ def fetch_input_variables(
                         upstream_in_dynamic_upstream = True
 
             # This is for blocks with multiple upstream dynamic blocks or dynamic child blocks.
-            if dynamic_block_indexes and \
-                    len(dynamic_block_indexes) >= 2 and \
-                    upstream_block_uuid in dynamic_block_indexes:
-
+            if (
+                dynamic_block_indexes
+                and len(dynamic_block_indexes) >= 2
+                and upstream_block_uuid in dynamic_block_indexes
+            ):
                 input_value = None
 
                 input_data = variable_values[0]
@@ -612,7 +663,9 @@ def fetch_input_variables(
                         # SQL blocks will return a Pandas DataFrame
                         if type(arr) is pd.DataFrame:
                             val = arr.iloc[dynamic_block_index].to_dict()
-                        elif type(arr) is list and len(arr) >= 1 and dynamic_block_index < len(arr):
+                        elif (
+                            type(arr) is list and len(arr) >= 1 and dynamic_block_index < len(arr)
+                        ):
                             val = arr[dynamic_block_index]
 
                 input_vars[idx] = val
@@ -681,6 +734,17 @@ def fetch_input_variables(
                             input_args=input_args,
                             partition=execution_partition,
                             spark=spark,
+                            read_chunks=current_block.upstream_chunks(upstream_block_uuid)
+                            if current_block
+                            else None,
+                            read_batch_settings=current_block.upstream_batch_settings(
+                                upstream_block_uuid
+                            )
+                            if current_block
+                            else None,
+                            input_data_types=current_block.input_data_types(upstream_block_uuid)
+                            if current_block
+                            else None,
                         )
 
                         if type(variable_values) is list and len(variable_values) == 1:
@@ -694,10 +758,11 @@ def fetch_input_variables(
                             final_value.append(val)
 
                     # output_0 is the metadata for dynamic blocks
-                    if dynamic_block_index is not None and \
-                            upstream_is_dynamic and \
-                            len(variables) >= 2:
-
+                    if (
+                        dynamic_block_index is not None
+                        and upstream_is_dynamic
+                        and len(variables) >= 2
+                    ):
                         var = variables[1]
                         variable_values = pipeline.get_block_variable(
                             upstream_block_uuid,
@@ -707,15 +772,27 @@ def fetch_input_variables(
                             input_args=input_args,
                             partition=execution_partition,
                             spark=spark,
+                            read_chunks=current_block.upstream_chunks(upstream_block_uuid)
+                            if current_block
+                            else None,
+                            read_batch_settings=current_block.upstream_batch_settings(
+                                upstream_block_uuid
+                            )
+                            if current_block
+                            else None,
+                            input_data_types=current_block.input_data_types(upstream_block_uuid)
+                            if current_block
+                            else None,
                         )
                         if dynamic_block_index < len(variable_values):
                             val = variable_values[dynamic_block_index]
                             kwargs_vars.append(val)
 
-                if ((upstream_is_dynamic and dynamic_block_index is not None)
-                    or should_reduce) and \
-                        len(final_value) >= 1 and \
-                        all([type(v) is pd.DataFrame for v in final_value]):
+                if (
+                    ((upstream_is_dynamic and dynamic_block_index is not None) or should_reduce)
+                    and len(final_value) >= 1
+                    and all([type(v) is pd.DataFrame for v in final_value])
+                ):
                     final_value = pd.concat(final_value)
 
                 if not should_reduce:
@@ -728,9 +805,11 @@ def fetch_input_variables(
                         else:
                             final_value = final_value[dynamic_block_index]
 
-                    if type(final_value) is not pd.DataFrame and \
-                        type(final_value) is list and \
-                            len(final_value) == 1:
+                    if (
+                        type(final_value) is not pd.DataFrame
+                        and type(final_value) is list
+                        and len(final_value) == 1
+                    ):
                         final_value = final_value[0]
 
                 input_vars[idx] = final_value
@@ -740,9 +819,29 @@ def fetch_input_variables(
         input_vars=input_vars,
         kwargs_vars=kwargs_vars,
         upstream_block_uuids_final=upstream_block_uuids_final,
-        __uuid='output_variables'
+        __uuid='output_variables',
     )
 
+    if kwargs_vars:
+        kwargs_vars2 = []
+
+        remote_blocks_output = []
+        for kwargs in kwargs_vars:
+            for remote_block_dict in kwargs.get('remote_blocks', []):
+                # Global data products only need the remote block information, not the output
+                if current_block and BlockType.GLOBAL_DATA_PRODUCT == current_block.type:
+                    output = remote_block_dict
+                else:
+                    output = RemoteBlock.load(**remote_block_dict).get_outputs()
+                remote_blocks_output.append(output)
+
+        for kwargs in kwargs_vars:
+            if kwargs.get('remote_blocks'):
+                kwargs['remote_blocks'] = remote_blocks_output
+            kwargs_vars2.append(kwargs)
+
+        kwargs_vars = kwargs_vars2
+
     return input_vars, kwargs_vars, upstream_block_uuids_final
 
 
@@ -754,9 +853,7 @@ def serialize_output(
 ):
     if type(data) is pd.DataFrame:
         if csv_lines_only:
-            data = dict(
-                table=data.to_csv(header=True, index=False).strip('\n').split('\n')
-            )
+            data = dict(table=data.to_csv(header=True, index=False).strip('\n').split('\n'))
         else:
             row_count, column_count = data.shape
 
@@ -764,9 +861,7 @@ def serialize_output(
             data = dict(
                 sample_data=dict(
                     columns=columns_to_display,
-                    rows=json.loads(
-                        data[columns_to_display].to_json(orient='split')
-                    )['data']
+                    rows=json.loads(data[columns_to_display].to_json(orient='split'))['data'],
                 ),
                 shape=[row_count, column_count],
                 type=DataType.TABLE,
@@ -774,11 +869,11 @@ def serialize_output(
             )
     elif is_geo_dataframe(data):
         data = dict(
-            text_data=f'''Use the code in a scratchpad to get the output of the block:
+            text_data=f"""Use the code in a scratchpad to get the output of the block:
 
 from mage_ai.data_preparation.variable_manager import get_variable
 df = get_variable('{block.pipeline.uuid}', '{block.uuid}', 'df')
-''',
+""",
             type=DataType.TEXT,
             variable_uuid=variable_uuid,
         )
@@ -804,7 +899,7 @@ df = get_variable('{block.pipeline.uuid}', '{block.uuid}', 'df')
         data = dict(
             sample_data=dict(
                 columns=columns_to_display,
-                rows=json.loads(df[columns_to_display].to_json(orient='split'))['data']
+                rows=json.loads(df[columns_to_display].to_json(orient='split'))['data'],
             ),
             type=DataType.TABLE,
             variable_uuid=variable_uuid,

mage_ai/data_preparation/models/constants.py

@@ -1,22 +1,26 @@
 import os
 from enum import Enum
 
-DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json'
-DATAFRAME_ANALYSIS_KEYS = frozenset(
-    [
-        'metadata',
-        'statistics',
-        'insights',
-        'suggestions',
-    ]
+from mage_ai.data_preparation.models.variables.constants import (
+    VariableAggregateDataType,
 )
+
+DATAFRAME_ANALYSIS_KEYS = frozenset([
+    VariableAggregateDataType.INSIGHTS.value,
+    VariableAggregateDataType.METADATA.value,
+    VariableAggregateDataType.STATISTICS.value,
+    VariableAggregateDataType.SUGGESTIONS.value,
+])
+DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json'
 DATAFRAME_ANALYSIS_MAX_COLUMNS = 100
 DATAFRAME_ANALYSIS_MAX_ROWS = 100_000
 DATAFRAME_SAMPLE_COUNT = 1000
 DATAFRAME_SAMPLE_COUNT_PREVIEW = 10
 DATAFRAME_SAMPLE_MAX_COLUMNS = 1000
-LOGS_DIR = '.logs'
+DYNAMIC_CHILD_BLOCK_SAMPLE_COUNT_PREVIEW = 10
+LOGS_SUBDIR = '.logs'
 MAX_PRINT_OUTPUT_LINES = int(os.getenv('MAX_PRINT_OUTPUT_LINES', 1000) or 1000)
+MAX_RESULTS_FOR_BLOCK_OUTPUTS_PREVIEW = 10
 PIPELINE_CONFIG_FILE = 'metadata.yaml'
 PIPELINE_MAX_FILE_SIZE = 500000  # maximum size of a pipeline import zip in kb (500Mb)
 PIPELINES_FOLDER = 'pipelines'
@@ -24,6 +28,8 @@ PREFERENCES_FILE = '.preferences.yaml'
 REPO_CONFIG_FILE = 'metadata.yaml'
 VARIABLE_DIR = '.variables'
 
+PIPELINE_RUN_STATUS_LAST_RUN_FAILED = 'last_run_failed'
+
 
 class AIMode(str, Enum):
     OPEN_AI = 'open_ai'
@@ -101,9 +107,9 @@ class PipelineType(str, Enum):
 
 
 class PipelineStatus(str, Enum):
-    ACTIVE = 'active',              # At least one active trigger
-    INACTIVE = 'inactive',          # All inactive triggers
-    NO_SCHEDULES = 'no_schedules',  # No triggers
+    ACTIVE = ('active',)  # At least one active trigger
+    INACTIVE = ('inactive',)  # All inactive triggers
+    NO_SCHEDULES = ('no_schedules',)  # No triggers
 
 
 BLOCK_LANGUAGE_TO_FILE_EXTENSION = {
@@ -155,8 +161,9 @@ SINGULAR_FOLDER_BLOCK_TYPES = [
     BlockType.CUSTOM,
 ]
 
-BLOCK_TYPE_DIRECTORY_NAME = \
-    {v: f'{v.value}s' for v in BlockType if v not in SINGULAR_FOLDER_BLOCK_TYPES}
+BLOCK_TYPE_DIRECTORY_NAME = {
+    v: f'{v.value}s' for v in BlockType if v not in SINGULAR_FOLDER_BLOCK_TYPES
+}
 
 
 PIPELINE_TYPE_DISPLAY_NAME_MAPPING = {