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

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
@@ -126,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:
@@ -143,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(
@@ -196,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.
@@ -222,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(
@@ -239,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,
@@ -420,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:
@@ -434,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)
@@ -455,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,
                         )
@@ -493,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
@@ -533,28 +548,34 @@ def fetch_input_variables(
             # 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:
@@ -571,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
                     ]
@@ -584,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]
@@ -630,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
@@ -699,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:
@@ -712,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,
@@ -725,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:
@@ -746,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
@@ -758,7 +819,7 @@ 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:
@@ -772,7 +833,7 @@ def fetch_input_variables(
                     output = remote_block_dict
                 else:
                     output = RemoteBlock.load(**remote_block_dict).get_outputs()
-            remote_blocks_output.append(output)
+                remote_blocks_output.append(output)
 
         for kwargs in kwargs_vars:
             if kwargs.get('remote_blocks'):
@@ -792,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
 
@@ -802,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,
@@ -812,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,
         )
@@ -842,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'
@@ -103,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 = {
@@ -157,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 = {

mage_ai/data_preparation/models/custom_templates/custom_block_template.py

@@ -1,8 +1,11 @@
 import os
 import shutil
-import yaml
 from dataclasses import dataclass, field
+from typing import Dict, List
+
+import yaml
 from jinja2 import Template
+
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.constants import (
     BLOCK_LANGUAGE_TO_FILE_EXTENSION,
@@ -10,18 +13,18 @@ from mage_ai.data_preparation.models.constants import (
     BlockLanguage,
     BlockType,
 )
-from mage_ai.data_preparation.models.file import File
-from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.custom_templates.constants import (
     DIRECTORY_FOR_BLOCK_TEMPLATES,
     METADATA_FILENAME_WITH_EXTENSION,
 )
-from mage_ai.data_preparation.models.custom_templates.utils import custom_templates_directory
-from mage_ai.settings.repo import get_repo_path
+from mage_ai.data_preparation.models.custom_templates.utils import (
+    custom_templates_directory,
+)
+from mage_ai.data_preparation.models.file import File
+from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.shared.config import BaseConfig
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.io import safe_write
-from typing import Dict, List
 
 
 @dataclass
@@ -34,12 +37,13 @@ class CustomBlockTemplate(BaseConfig):
     language: BlockLanguage = None
     name: str = None
     pipeline: Dict = field(default_factory=dict)
+    repo_path: str = None
     tags: List = field(default_factory=list)
     template_uuid: str = None
     user: Dict = field(default_factory=dict)
 
     @classmethod
-    def load(self, template_uuid: str = None, uuid: str = None):
+    def load(self, repo_path, template_uuid: str = None, uuid: str = None):
         uuid_use = uuid
         template_uuid_use = template_uuid
 
@@ -55,12 +59,13 @@ class CustomBlockTemplate(BaseConfig):
 
         try:
             config_path_metadata = os.path.join(
-                get_repo_path(),
+                repo_path,
                 uuid_use,
                 METADATA_FILENAME_WITH_EXTENSION,
             )
             custom_template = super().load(config_path_metadata)
             custom_template.template_uuid = template_uuid_use
+            custom_template.repo_path = repo_path
 
             return custom_template
         except Exception as err:
@@ -77,7 +82,7 @@ class CustomBlockTemplate(BaseConfig):
     @property
     def metadata_file_path(self) -> str:
         return os.path.join(
-            get_repo_path(),
+            self.repo_path,
             self.uuid,
             METADATA_FILENAME_WITH_EXTENSION,
         )
@@ -98,7 +103,7 @@ class CustomBlockTemplate(BaseConfig):
         return Block.create(
             block_name,
             self.block_type,
-            get_repo_path(),
+            self.repo_path,
             color=self.color,
             configuration=configuration,
             extension_uuid=extension_uuid,
@@ -118,7 +123,7 @@ class CustomBlockTemplate(BaseConfig):
         return File(
             dir_path=self.uuid,
             filename=filename,
-            repo_path=get_repo_path(),
+            repo_path=self.repo_path,
         ).content()
 
     def render_template(
@@ -170,8 +175,8 @@ class CustomBlockTemplate(BaseConfig):
                 filename,
                 self.uuid,
                 self.content,
-                get_repo_path(),
+                self.repo_path,
             )
 
     def delete(self) -> None:
-        shutil.rmtree(os.path.join(get_repo_path(), self.uuid))
+        shutil.rmtree(os.path.join(self.repo_path, self.uuid))