mage-ai 0.8.91__tar.gz → 0.8.93__tar.gz

{mage-ai-0.8.91/mage_ai.egg-info → mage-ai-0.8.93}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.8.91
+Version: 0.8.93
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/api/policies/BlockPolicy.py

@@ -81,6 +81,7 @@ BlockPolicy.allow_write([
     'bookmark_values',
     'callback_blocks',
     'color',
+    'conditional_blocks',
     'configuration',
     'content',
     'destination_table',

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/api/policies/PipelinePolicy.py

@@ -37,6 +37,7 @@ PipelinePolicy.allow_actions([
 
 PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
     'callbacks',
+    'conditionals',
     'extensions',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
@@ -46,6 +47,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
 
 PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
     'callbacks',
+    'conditionals',
     'extensions',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
@@ -56,6 +58,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
 
 PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
     'callbacks',
+    'conditionals',
     'extensions',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
@@ -65,6 +68,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
 
 PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
     'callbacks',
+    'conditionals',
     'extensions',
     'schedules',
 ], scopes=[
@@ -75,6 +79,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
 
 PipelinePolicy.allow_write([
     'callbacks',
+    'conditionals',
     'clone_pipeline_uuid',
     'extensions',
     'name',
@@ -88,6 +93,7 @@ PipelinePolicy.allow_write([
 PipelinePolicy.allow_write([
     'add_upstream_for_block_uuid',
     'callbacks',
+    'conditionals',
     'extensions',
     'schedules',
 ] + PipelinePresenter.default_attributes, scopes=[

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/api/presenters/BlockPresenter.py

@@ -11,6 +11,7 @@ class BlockPresenter(BasePresenter):
         'all_upstream_blocks_executed',
         'callback_blocks',
         'color',
+        'conditional_blocks',
         'configuration',
         'downstream_blocks',
         'executor_config',

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/api/presenters/PipelinePresenter.py

@@ -12,6 +12,7 @@ class PipelinePresenter(BasePresenter):
         'executor_count',
         'executor_type',
         'name',
+        'spark_config',
         'type',
         'updated_at',
         'uuid',
@@ -44,6 +45,7 @@ class PipelinePresenter(BasePresenter):
                 include_block_metadata=include_block_metadata,
                 include_block_tags=True,
                 include_callback_blocks=True,
+                include_conditional_blocks=True,
                 include_content=include_content,
                 include_extensions=include_extensions,
                 include_outputs=include_outputs,

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/data_preparation/decorators.py

@@ -9,6 +9,10 @@ def callback(function):
     return function
 
 
+def condition(function):
+    return function
+
+
 def custom(function):
     return function
 

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/data_preparation/executors/block_executor.py

@@ -64,6 +64,31 @@ class BlockExecutor:
                 on_start(self.block_uuid)
             pipeline_run = PipelineRun.query.get(kwargs['pipeline_run_id']) \
                 if 'pipeline_run_id' in kwargs else None
+
+            conditional_result = self._execute_conditional(
+                dynamic_block_index=dynamic_block_index,
+                dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                global_vars=global_vars,
+                logging_tags=tags,
+                pipeline_run=pipeline_run,
+            )
+            if not conditional_result:
+                self.logger.info(
+                    f'Conditional block(s) returned false for {self.block.uuid}. '
+                    'This block run and downstream blocks will set as CONDITION_FAILED.',
+                    **merge_dict(tags, dict(
+                        block_type=self.block.type,
+                        block_uuid=self.block.uuid,
+                    )),
+                )
+                self.__update_block_run_status(
+                    'condition_failed',
+                    block_run_id=kwargs.get('block_run_id'),
+                    callback_url=callback_url,
+                    tags=tags,
+                )
+                return dict(output=[])
+
             try:
                 from mage_ai.shared.retry import retry
 
@@ -207,12 +232,52 @@ class BlockExecutor:
 
         return result
 
+    def _execute_conditional(
+        self,
+        global_vars: Dict,
+        logging_tags: Dict,
+        pipeline_run: PipelineRun,
+        dynamic_block_index: Union[int, None] = None,
+        dynamic_upstream_block_uuids: Union[List[str], None] = None,
+    ):
+        result = True
+        for conditional_block in self.block.conditional_blocks:
+            try:
+                block_result = conditional_block.execute_conditional(
+                    self.block,
+                    dynamic_block_index=dynamic_block_index,
+                    dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                    execution_partition=self.execution_partition,
+                    global_vars=global_vars,
+                    logger=self.logger,
+                    logging_tags=logging_tags,
+                    pipeline_run=pipeline_run,
+                )
+                if not block_result:
+                    self.logger.info(
+                        f'Conditional block {conditional_block.uuid} evaluated as False ' +
+                        f'for block {self.block.uuid}',
+                        logging_tags,
+                    )
+                result = result and block_result
+            except Exception as conditional_err:
+                self.logger.exception(
+                    f'Failed to execute conditional block {conditional_block.uuid} ' +
+                    f'for block {self.block.uuid}.',
+                    **merge_dict(logging_tags, dict(
+                        error=conditional_err,
+                    )),
+                )
+                result = False
+
+        return result
+
     def _execute_callback(
         self,
         callback: str,
-        global_vars,
-        logging_tags,
-        pipeline_run,
+        global_vars: Dict,
+        logging_tags: Dict,
+        pipeline_run: PipelineRun,
         dynamic_block_index: Union[int, None] = None,
         dynamic_upstream_block_uuids: Union[List[str], None] = None,
     ):

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/data_preparation/models/block/__init__.py

@@ -47,7 +47,7 @@ from mage_ai.data_preparation.models.constants import (
 )
 from mage_ai.data_preparation.models.file import File
 from mage_ai.data_preparation.models.variable import VariableType
-from mage_ai.data_preparation.repo_manager import get_repo_path
+from mage_ai.data_preparation.repo_manager import RepoConfig, get_repo_path
 from mage_ai.data_preparation.shared.stream import StreamToLogger
 from mage_ai.data_preparation.templates.template import load_template
 from mage_ai.server.kernel_output_parser import DataType
@@ -253,6 +253,7 @@ class Block:
 
         self._outputs = None
         self._outputs_loaded = False
+        self.conditional_blocks = []
         self.callback_blocks = []
         self.upstream_blocks = []
         self.downstream_blocks = []
@@ -348,6 +349,10 @@ class Block:
     def callback_block_uuids(self) -> List[str]:
         return [b.uuid for b in self.callback_blocks]
 
+    @property
+    def conditional_block_uuids(self) -> List[str]:
+        return [b.uuid for b in self.conditional_blocks]
+
     @property
     def upstream_block_uuids(self) -> List[str]:
         return [b.uuid for b in self.upstream_blocks]
@@ -456,6 +461,8 @@ class Block:
         elif pipeline and PipelineType.INTEGRATION == pipeline.type:
             if BlockType.CALLBACK == block_type:
                 return CallbackBlock
+            elif BlockType.CONDITIONAL == block_type:
+                return ConditionalBlock
             elif BlockType.DATA_LOADER == block_type:
                 return SourceBlock
             elif BlockType.DATA_EXPORTER == block_type:
@@ -670,11 +677,33 @@ class Block:
         if logging_tags is None:
             logging_tags = dict()
 
-        arr = []
+        if self.conditional_blocks and len(self.conditional_blocks) > 0:
+            conditional_message = ''
+            result = True
+            for conditional_block in self.conditional_blocks:
+                block_result = conditional_block.execute_conditional(
+                    self,
+                    global_vars=global_vars,
+                    logger=logger,
+                    logging_tags=logging_tags,
+                )
+                conditional_message += \
+                    f'Conditional block {conditional_block.uuid} evaluated to {block_result}.\n'
+                result = result and block_result
+
+            # Print result to block output
+            if not result:
+                conditional_message += 'This block would not be executed in a trigger run.\n'
+            conditional_json = json.dumps(dict(
+                message=conditional_message,
+            ))
+            print(f'[__internal_test__]{conditional_json}')
+
+        callback_arr = []
         if self.callback_block:
-            arr.append(self.callback_block)
+            callback_arr.append(self.callback_block)
         if self.callback_blocks:
-            arr += self.callback_blocks
+            callback_arr += self.callback_blocks
 
         try:
             output = self.execute_sync(
@@ -684,7 +713,7 @@ class Block:
                 **kwargs
             )
         except Exception as e:
-            for callback_block in arr:
+            for callback_block in callback_arr:
                 callback_block.execute_callback(
                     'on_failure',
                     global_vars=global_vars,
@@ -694,7 +723,7 @@ class Block:
                 )
             raise e
 
-        for callback_block in arr:
+        for callback_block in callback_arr:
             callback_block.execute_callback(
                 'on_success',
                 global_vars=global_vars,
@@ -1131,6 +1160,8 @@ class Block:
         self,
         execution_partition: str = None,
         include_print_outputs: bool = True,
+        csv_lines_only: bool = False,
+        sample: bool = True,
         sample_count: int = DATAFRAME_SAMPLE_COUNT_PREVIEW,
         variable_type: VariableType = None,
         block_uuid: str = None,
@@ -1166,40 +1197,47 @@ class Block:
                 continue
 
             data = variable_object.read_data(
-                sample=True,
+                sample=sample,
                 sample_count=sample_count,
                 spark=self.__get_spark_session(),
             )
             if type(data) is pd.DataFrame:
-                try:
-                    analysis = variable_manager.get_variable(
-                        self.pipeline.uuid,
-                        block_uuid,
-                        v,
-                        dataframe_analysis_keys=['metadata', 'statistics'],
-                        partition=execution_partition,
-                        variable_type=VariableType.DATAFRAME_ANALYSIS,
+                if csv_lines_only:
+                    data = dict(
+                        table=data.to_csv(header=True, index=False).strip('\n').split('\n')
                     )
-                except Exception:
-                    analysis = None
-                if analysis is not None:
-                    stats = analysis.get('statistics', {})
-                    column_types = (analysis.get('metadata') or {}).get('column_types', {})
-                    row_count = stats.get('original_row_count', stats.get('count'))
-                    column_count = stats.get('original_column_count', len(column_types))
                 else:
-                    row_count, column_count = data.shape
-
-                columns_to_display = data.columns.tolist()[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
-                data = dict(
-                    sample_data=dict(
-                        columns=columns_to_display,
-                        rows=json.loads(data[columns_to_display].to_json(orient='split'))['data']
-                    ),
-                    shape=[row_count, column_count],
-                    type=DataType.TABLE,
-                    variable_uuid=v,
-                )
+                    try:
+                        analysis = variable_manager.get_variable(
+                            self.pipeline.uuid,
+                            block_uuid,
+                            v,
+                            dataframe_analysis_keys=['metadata', 'statistics'],
+                            partition=execution_partition,
+                            variable_type=VariableType.DATAFRAME_ANALYSIS,
+                        )
+                    except Exception:
+                        analysis = None
+                    if analysis is not None:
+                        stats = analysis.get('statistics', {})
+                        column_types = (analysis.get('metadata') or {}).get('column_types', {})
+                        row_count = stats.get('original_row_count', stats.get('count'))
+                        column_count = stats.get('original_column_count', len(column_types))
+                    else:
+                        row_count, column_count = data.shape
+
+                    columns_to_display = data.columns.tolist()[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
+                    data = dict(
+                        sample_data=dict(
+                            columns=columns_to_display,
+                            rows=json.loads(
+                                data[columns_to_display].to_json(orient='split')
+                            )['data']
+                        ),
+                        shape=[row_count, column_count],
+                        type=DataType.TABLE,
+                        variable_uuid=v,
+                    )
                 data_products.append(data)
                 continue
             elif is_geo_dataframe(data):
@@ -1383,7 +1421,11 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         variable_mapping = self.__save_outputs_prepare(outputs)
         await self.store_variables_async(variable_mapping, override=override)
 
-    def to_dict_base(self, include_callback_blocks: bool = False) -> Dict:
+    def to_dict_base(
+        self,
+        include_callback_blocks: bool = False,
+        include_conditional_blocks: bool = False,
+    ) -> Dict:
         language = self.language
         if language and type(self.language) is not str:
             language = self.language.value
@@ -1410,6 +1452,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         if include_callback_blocks:
             data['callback_blocks'] = self.callback_block_uuids
 
+        if include_conditional_blocks:
+            data['conditional_blocks'] = self.conditional_block_uuids
+
         if self.replicated_block:
             data['replicated_block'] = self.replicated_block
 
@@ -1451,12 +1496,16 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         include_block_metadata: bool = False,
         include_block_tags: bool = False,
         include_callback_blocks: bool = False,
+        include_conditional_blocks: bool = False,
         include_content: bool = False,
         include_outputs: bool = False,
         sample_count: int = None,
         check_if_file_exists: bool = False,
     ) -> Dict:
-        data = self.to_dict_base(include_callback_blocks=include_callback_blocks)
+        data = self.to_dict_base(
+            include_callback_blocks=include_callback_blocks,
+            include_conditional_blocks=include_conditional_blocks,
+        )
 
         if include_content:
             data['content'] = await self.content_async()
@@ -1509,6 +1558,11 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         ):
             self.__update_callback_blocks(data['callback_blocks'])
 
+        if 'conditional_blocks' in data and set(data['conditional_blocks']) != set(
+            self.conditional_block_uuids
+        ):
+            self.__update_conditional_blocks(data['conditional_blocks'])
+
         if 'executor_type' in data and data['executor_type'] != self.executor_type:
             self.executor_type = data['executor_type']
             self.__update_pipeline_block()
@@ -1524,6 +1578,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
     def update_callback_blocks(self, callback_blocks: List[Any]) -> None:
         self.callback_blocks = callback_blocks
 
+    def update_conditional_blocks(self, conditional_blocks: List[Any]) -> None:
+        self.conditional_blocks = conditional_blocks
+
     def update_upstream_blocks(self, upstream_blocks: List[Any]) -> None:
         self.upstream_blocks = upstream_blocks
 
@@ -1839,11 +1896,17 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         return global_vars
 
     def __get_spark_session(self):
-        if self.spark_init:
+        if self.spark_init and (not self.pipeline or
+                                not self.pipeline.spark_config):
             return self.spark
         try:
-            spark_config = SparkConfig.load(
-                config={'repo_path': self.repo_path})
+            if self.pipeline and self.pipeline.spark_config:
+                spark_config = SparkConfig.load(
+                    config=self.pipeline.spark_config)
+            else:
+                repo_config = RepoConfig(repo_path=self.repo_path)
+                spark_config = SparkConfig.load(
+                    config=repo_config.spark_config)
             self.spark = get_spark_session(spark_config)
         except Exception:
             self.spark = None
@@ -2026,6 +2089,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
 
     def tags(self) -> List[str]:
         from mage_ai.data_preparation.models.block.constants import (
+            TAG_CONDITION,
             TAG_DYNAMIC,
             TAG_DYNAMIC_CHILD,
             TAG_REDUCE_OUTPUT,
@@ -2046,6 +2110,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         if self.replicated_block:
             arr.append(TAG_REPLICA)
 
+        if len(self.conditional_blocks) > 0:
+            arr.append(TAG_CONDITION)
+
         return arr
 
     def variable_object(
@@ -2155,6 +2222,16 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
             widget=BlockType.CHART == self.type,
         )
 
+    def __update_conditional_blocks(self, block_uuids: List[str]) -> None:
+        if self.pipeline is None:
+            return
+
+        self.pipeline.update_block(
+            self,
+            conditional_block_uuids=block_uuids,
+            widget=BlockType.CHART == self.type,
+        )
+
 
 class SensorBlock(Block):
     def execute_block_function(
@@ -2187,7 +2264,98 @@ class SensorBlock(Block):
             return []
 
 
-class CallbackBlock(Block):
+class AddonBlock(Block):
+    def _create_global_vars(
+        self,
+        global_vars: Dict,
+        parent_block: Block,
+        **kwargs,
+    ) -> Dict:
+        pipeline_run = kwargs.get('pipeline_run')
+        global_vars = merge_dict(
+            global_vars or dict(),
+            dict(
+                pipeline_uuid=self.pipeline.uuid,
+                block_uuid=self.uuid,
+                pipeline_run=pipeline_run,
+            ),
+        )
+        if parent_block:
+            global_vars['parent_block_uuid'] = parent_block.uuid
+
+        if parent_block and \
+                parent_block.pipeline and \
+                PipelineType.INTEGRATION == parent_block.pipeline.type:
+
+            template_runtime_configuration = parent_block.template_runtime_configuration
+            index = template_runtime_configuration.get('index', None)
+            is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
+            selected_streams = template_runtime_configuration.get('selected_streams', [])
+            stream = selected_streams[0] if len(selected_streams) >= 1 else None
+            destination_table = template_runtime_configuration.get('destination_table', stream)
+
+            global_vars['index'] = index
+            global_vars['is_last_block_run'] = is_last_block_run
+            global_vars['stream'] = stream
+            global_vars['destination_table'] = destination_table
+
+        return global_vars
+
+
+class ConditionalBlock(AddonBlock):
+    def execute_conditional(
+        self,
+        parent_block: Block,
+        dynamic_block_index: Union[int, None] = None,
+        dynamic_upstream_block_uuids: Union[List[str], None] = None,
+        execution_partition: str = None,
+        global_vars: Dict = None,
+        logger: Logger = None,
+        logging_tags: Dict = None,
+        **kwargs
+    ) -> bool:
+        if logger is not None:
+            stdout = StreamToLogger(logger, logging_tags=logging_tags)
+        else:
+            stdout = sys.stdout
+
+        with redirect_stdout(stdout):
+            global_vars = self._create_global_vars(
+                global_vars,
+                parent_block,
+                **kwargs,
+            )
+
+            condition_functions = []
+
+            results = dict(
+                condition=self._block_decorator(condition_functions),
+            )
+            exec(self.content, results)
+
+            global_vars_copy = global_vars.copy()
+            input_vars = []
+            # Fetch input variables for parent block
+            if parent_block is not None:
+                input_vars, kwargs_vars, _ = parent_block.fetch_input_variables(
+                    None,
+                    execution_partition,
+                    global_vars,
+                    dynamic_block_index=dynamic_block_index,
+                    dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                )
+
+                for kwargs_var in kwargs_vars:
+                    global_vars_copy.update(kwargs_var)
+
+            result = True
+            for condition_function in condition_functions:
+                result = condition_function(*input_vars, **global_vars_copy) and result
+
+            return result
+
+
+class CallbackBlock(AddonBlock):
     @classmethod
     def create(cls, orig_block_name) -> 'CallbackBlock':
         return Block.create(
@@ -2209,40 +2377,17 @@ class CallbackBlock(Block):
         parent_block: Block = None,
         **kwargs
     ) -> None:
-        pipeline_run = kwargs.get('pipeline_run')
-
         if logger is not None:
             stdout = StreamToLogger(logger, logging_tags=logging_tags)
         else:
             stdout = sys.stdout
 
         with redirect_stdout(stdout):
-            global_vars = merge_dict(
-                global_vars or dict(),
-                dict(
-                    pipeline_uuid=self.pipeline.uuid,
-                    block_uuid=self.uuid,
-                    pipeline_run=pipeline_run,
-                ),
+            global_vars = self._create_global_vars(
+                global_vars,
+                parent_block,
+                **kwargs
             )
-            if parent_block:
-                global_vars['parent_block_uuid'] = parent_block.uuid
-
-            if parent_block and \
-                    parent_block.pipeline and \
-                    PipelineType.INTEGRATION == parent_block.pipeline.type:
-
-                template_runtime_configuration = parent_block.template_runtime_configuration
-                index = template_runtime_configuration.get('index', None)
-                is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
-                selected_streams = template_runtime_configuration.get('selected_streams', [])
-                stream = selected_streams[0] if len(selected_streams) >= 1 else None
-                destination_table = template_runtime_configuration.get('destination_table', stream)
-
-                global_vars['index'] = index
-                global_vars['is_last_block_run'] = is_last_block_run
-                global_vars['stream'] = stream
-                global_vars['destination_table'] = destination_table
 
             callback_functions = []
             failure_functions = []

{mage-ai-0.8.91 → mage-ai-0.8.93}/mage_ai/data_preparation/models/block/constants.py

@@ -1,15 +1,16 @@
 from mage_ai.data_preparation.models.block import (
     Block,
     CallbackBlock,
+    ConditionalBlock,
     SensorBlock,
 )
 from mage_ai.data_preparation.models.block.dbt import DBTBlock
 from mage_ai.data_preparation.models.block.extension.block import ExtensionBlock
 from mage_ai.data_preparation.models.constants import BlockType
 
-
 BLOCK_TYPE_TO_CLASS = {
     BlockType.CALLBACK: CallbackBlock,
+    BlockType.CONDITIONAL: ConditionalBlock,
     BlockType.CUSTOM: Block,
     BlockType.DATA_EXPORTER: Block,
     BlockType.DATA_LOADER: Block,
@@ -21,6 +22,7 @@ BLOCK_TYPE_TO_CLASS = {
     BlockType.SENSOR: SensorBlock,
 }
 
+TAG_CONDITION = 'condition'
 TAG_DYNAMIC = 'dynamic'
 TAG_DYNAMIC_CHILD = 'dynamic_child'
 TAG_REDUCE_OUTPUT = 'reduce_output'