mage-ai 0.9.46__py3-none-any.whl → 0.9.48__py3-none-any.whl

mage_ai/data_integrations/utils/scheduler.py

@@ -8,7 +8,9 @@ from mage_ai.data_integrations.sources.constants import SQL_SOURCES
 from mage_ai.data_integrations.utils.config import build_config, get_batch_fetch_limit
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.models.block.data_integration.constants import (
+    KEY_REPLICATION_METHOD,
     MAX_QUERY_STRING_SIZE,
+    REPLICATION_METHOD_INCREMENTAL,
 )
 from mage_ai.data_preparation.models.block.data_integration.utils import (
     convert_block_output_data_for_destination,
@@ -19,6 +21,7 @@ from mage_ai.data_preparation.models.block.data_integration.utils import (
 from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
     IntegrationPipeline,
 )
+from mage_ai.data_preparation.models.triggers import ScheduleInterval
 from mage_ai.orchestration.db import db_connection
 from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun
 from mage_ai.orchestration.metrics.pipeline_run import calculate_metrics
@@ -253,6 +256,7 @@ def build_block_run_metadata(
     logging_tags: Dict = None,
     parent_stream: str = None,
     partition: str = None,
+    pipeline_run: PipelineRun = None,
     selected_streams: List[str] = None,
 ) -> List[Dict]:
     block_run_metadata = []
@@ -270,9 +274,11 @@ def build_block_run_metadata(
 
     if block.is_source():
         return __build_block_run_metadata_for_source(
+            block,
             data_integration_settings,
             logger,
             logging_tags=logging_tags,
+            pipeline_run=pipeline_run,
             selected_streams=selected_streams,
         )
 
@@ -377,9 +383,11 @@ def __build_block_run_metadata_for_destination(
 
 
 def __build_block_run_metadata_for_source(
+    block,
     data_integration_settings: Dict,
     logger: DictLogger,
     logging_tags: Dict = None,
+    pipeline_run: PipelineRun = None,
     selected_streams: List[str] = None,
 ) -> List[Dict]:
     block_run_metadata = []
@@ -388,20 +396,66 @@ def __build_block_run_metadata_for_source(
     config = data_integration_settings.get('config')
     batch_fetch_limit = get_batch_fetch_limit(config)
 
-    streams = selected_streams or \
-        [s.get('tap_stream_id') for s in get_selected_streams(catalog)]
+    stream_dicts_by_stream_id = index_by(
+        lambda x: x.get('tap_stream_id') or x.get('stream'),
+        get_selected_streams(catalog),
+    )
+
+    streams = []
+
+    if selected_streams:
+        streams = selected_streams
+    else:
+        streams = list(stream_dicts_by_stream_id.keys())
+
+    at_least_one_incremental = False
+
+    for stream_id in streams:
+        if at_least_one_incremental:
+            break
+
+        stream_dict = stream_dicts_by_stream_id.get(stream_id)
+        if not stream_dict:
+            continue
+
+        if REPLICATION_METHOD_INCREMENTAL == stream_dict.get(KEY_REPLICATION_METHOD):
+            at_least_one_incremental = True
+            break
+
+    execution_partition_previous = None
+
+    if at_least_one_incremental and pipeline_run:
+        pipeline_runs_completed = \
+            PipelineRun.recently_completed_pipeline_runs(
+                pipeline_run.pipeline_uuid,
+                pipeline_run_id=pipeline_run.id,
+                pipeline_schedule_id=(
+                    None if
+                    ScheduleInterval.ONCE == pipeline_run.pipeline_schedule.schedule_interval else
+                    pipeline_run.pipeline_schedule_id
+                ),
+                sample_size=1,
+            )
+
+        if pipeline_runs_completed:
+            execution_partition_previous = pipeline_runs_completed[0].execution_partition
+
     data_integration_uuid = data_integration_settings.get('data_integration_uuid')
 
     is_sql_source = data_integration_uuid in SQL_SOURCES_UUID
     record_counts_by_stream = {}
     if is_sql_source:
+
         record_counts_by_stream = index_by(
             lambda x: x['id'],
             count_records(
                 config,
                 data_integration_uuid,
                 streams,
+                block=block,
                 catalog=catalog,
+                partition=execution_partition_previous,
+                variables=pipeline_run.variables if pipeline_run else None,
             ),
         )
 
@@ -432,6 +486,7 @@ def __build_block_run_metadata_for_source(
 
         for idx in range(number_of_batches):
             block_run_metadata.append(dict(
+                execution_partition_previous=execution_partition_previous,
                 index=idx,
                 number_of_batches=number_of_batches,
                 stream=tap_stream_id,

mage_ai/data_preparation/executors/block_executor.py

@@ -84,6 +84,8 @@ class BlockExecutor:
         self,
         analyze_outputs: bool = False,
         block_run_id: int = None,
+        block_run_outputs_cache: Dict[str, List] = None,
+        cache_block_output_in_memory: bool = False,
         callback_url: Union[str, None] = None,
         global_vars: Union[Dict, None] = None,
         input_from_output: Union[Dict, None] = None,
@@ -105,6 +107,10 @@ class BlockExecutor:
         Args:
             analyze_outputs: Whether to analyze the outputs of the block.
             block_run_id: The ID of the block run.
+            block_run_outputs_cache: block uuid to block outputs mapping. It's used when
+                cache_block_output_in_memory is set to True.
+            cache_block_output_in_memory: Whether to cache the block output in memory. By default,
+                the block output is persisted on disk.
             callback_url: The URL for the callback.
             global_vars: Global variables for the block execution.
             input_from_output: Input from the output of a previous block.
@@ -512,6 +518,8 @@ class BlockExecutor:
                         return self._execute(
                             analyze_outputs=analyze_outputs,
                             block_run_id=block_run_id,
+                            block_run_outputs_cache=block_run_outputs_cache,
+                            cache_block_output_in_memory=cache_block_output_in_memory,
                             callback_url=callback_url,
                             global_vars=global_vars,
                             update_status=update_status,
@@ -620,6 +628,8 @@ class BlockExecutor:
         self,
         analyze_outputs: bool = False,
         block_run_id: int = None,
+        block_run_outputs_cache: Dict[str, List] = None,
+        cache_block_output_in_memory: bool = False,
         callback_url: Union[str, None] = None,
         global_vars: Union[Dict, None] = None,
         update_status: bool = False,
@@ -640,7 +650,11 @@ class BlockExecutor:
 
         Args:
             analyze_outputs: Whether to analyze the outputs of the block.
+            block_run_outputs_cache: block uuid to block outputs mapping. It's used when
+                cache_block_output_in_memory is set to True.
             callback_url: The URL for the callback.
+            cache_block_output_in_memory: Whether to cache the block output in memory. By default,
+                the block output is persisted on disk.
             global_vars: Global variables for the block execution.
             update_status: Whether to update the status of the block execution.
             input_from_output: Input from the output of a previous block.
@@ -659,7 +673,10 @@ class BlockExecutor:
             logging_tags = dict()
 
         extra_options = {}
-        store_variables = True
+        if cache_block_output_in_memory:
+            store_variables = False
+        else:
+            store_variables = True
         is_data_integration = False
 
         if self.project.is_feature_enabled(FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE):
@@ -735,6 +752,14 @@ class BlockExecutor:
                 except Exception as err:
                     print(f'[WARNING] BlockExecutor._execute: {err}')
 
+            is_source = self.block.is_source()
+            if is_source and data_integration_metadata:
+                execution_partition_previous = data_integration_metadata.get(
+                    'execution_partition_previous',
+                )
+                if execution_partition_previous:
+                    extra_options['execution_partition_previous'] = execution_partition_previous
+
             if di_settings and \
                     data_integration_metadata and \
                     data_integration_metadata.get('controller') and \
@@ -746,7 +771,6 @@ class BlockExecutor:
                 if is_data_integration:
                     arr = []
 
-                    is_source = self.block.is_source()
                     data_integration_uuid = di_settings.get('data_integration_uuid')
                     catalog = di_settings.get('catalog', [])
 
@@ -770,6 +794,7 @@ class BlockExecutor:
                             logging_tags=logging_tags,
                             parent_stream=data_integration_metadata.get('parent_stream'),
                             partition=self.execution_partition,
+                            pipeline_run=pipeline_run,
                             selected_streams=[stream],
                         )
                         for br_metadata in block_run_metadata:
@@ -949,6 +974,7 @@ class BlockExecutor:
 
         result = self.block.execute_sync(
             analyze_outputs=analyze_outputs,
+            block_run_outputs_cache=block_run_outputs_cache,
             execution_partition=self.execution_partition,
             global_vars=global_vars,
             logger=self.logger,
@@ -971,6 +997,7 @@ class BlockExecutor:
                 global_vars=global_vars,
                 logger=self.logger,
                 logging_tags=logging_tags,
+                outputs=result if cache_block_output_in_memory else None,
             )
         elif PipelineType.INTEGRATION != self.pipeline.type and \
                 (not is_data_integration or BlockLanguage.PYTHON == self.block.language):
@@ -979,6 +1006,7 @@ class BlockExecutor:
                 global_vars=global_vars,
                 logger=self.logger,
                 logging_tags=logging_tags,
+                outputs=result if cache_block_output_in_memory else None,
                 update_tests=False,
                 dynamic_block_uuid=dynamic_block_uuid,
             )

mage_ai/data_preparation/executors/pipeline_executor.py

@@ -89,7 +89,10 @@ class PipelineExecutor:
         if global_vars is None:
             global_vars = dict()
 
-        def create_block_task(block_run: BlockRun) -> asyncio.Task:
+        def create_block_task(
+            block_run: BlockRun,
+            block_run_outputs_cache: Dict[str, List],
+        ) -> asyncio.Task:
             async def execute_block() -> None:
                 executor_kwargs = dict(
                     pipeline=self.pipeline,
@@ -102,22 +105,34 @@ class PipelineExecutor:
                     block_run_data['started_at'] = datetime.now(tz=pytz.UTC)
 
                 block_run.update(**block_run_data)
-                BlockExecutor(**executor_kwargs).execute(
+                return BlockExecutor(**executor_kwargs).execute(
                     block_run_id=block_run.id,
+                    block_run_outputs_cache=block_run_outputs_cache,
+                    cache_block_output_in_memory=self.pipeline.cache_block_output_in_memory,
                     global_vars=global_vars,
                     pipeline_run_id=pipeline_run.id,
                 )
 
             return asyncio.create_task(execute_block())
 
+        block_run_outputs_cache = dict()
+
         while not pipeline_run.all_blocks_completed(allow_blocks_to_fail):
+            # Update the statuses of the block runs to CONDITION_FAILED or UPSTREAM_FAILED.
+            pipeline_run.update_block_run_statuses(pipeline_run.initial_block_runs)
             executable_block_runs = pipeline_run.executable_block_runs(
                 allow_blocks_to_fail=allow_blocks_to_fail,
             )
             if not executable_block_runs:
                 return
-            block_run_tasks = [create_block_task(b) for b in executable_block_runs]
-            await asyncio.gather(*block_run_tasks)
+            block_run_tasks = [
+                create_block_task(b, block_run_outputs_cache=block_run_outputs_cache)
+                for b in executable_block_runs]
+            block_run_outputs = await asyncio.gather(*block_run_tasks)
+            if self.pipeline.cache_block_output_in_memory:
+                for idx, block_run in enumerate(executable_block_runs):
+                    block_run_outputs_cache[block_run.block_uuid] = \
+                        block_run_outputs[idx].get('output', [])
 
     def build_tags(self, **kwargs):
         default_tags = dict(

mage_ai/data_preparation/models/block/__init__.py

@@ -23,6 +23,7 @@ from jinja2 import Template
 import mage_ai.data_preparation.decorators
 from mage_ai.cache.block import BlockCache
 from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
+from mage_ai.data_integrations.sources.constants import SQL_SOURCES_MAPPING
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.logging.logger_manager_factory import LoggerManagerFactory
 from mage_ai.data_preparation.models.block.data_integration.mixins import (
@@ -466,11 +467,14 @@ class Block(DataIntegrationMixin, SparkBlock):
                     uuid = settings.get('source') or settings.get('destination')
                     mapping = grouped_templates.get(uuid) or {}
 
+                    di_metadata = merge_dict(
+                        extract(mapping or {}, ['name']),
+                        settings,
+                    )
+                    di_metadata['sql'] = uuid in SQL_SOURCES_MAPPING
+
                     return dict(
-                        data_integration=merge_dict(
-                            extract(mapping or {}, ['name']),
-                            settings,
-                        ),
+                        data_integration=di_metadata,
                     )
             elif BlockLanguage.PYTHON == self.language:
                 try:
@@ -482,15 +486,18 @@ class Block(DataIntegrationMixin, SparkBlock):
                     uuid = di_settings.get('data_integration_uuid')
                     mapping = grouped_templates.get(uuid) or {}
 
+                    di_metadata = merge_dict(
+                        extract(mapping or {}, ['name']),
+                        ignore_keys(di_settings or {}, [
+                            'catalog',
+                            'config',
+                            'data_integration_uuid',
+                        ]),
+                    )
+                    di_metadata['sql'] = uuid in SQL_SOURCES_MAPPING
+
                     return dict(
-                        data_integration=merge_dict(
-                            extract(mapping or {}, ['name']),
-                            ignore_keys(di_settings or {}, [
-                                'catalog',
-                                'config',
-                                'data_integration_uuid',
-                            ]),
-                        ),
+                        data_integration=di_metadata,
                     )
                 except Exception as err:
                     if is_debug():
@@ -971,6 +978,7 @@ class Block(DataIntegrationMixin, SparkBlock):
     def execute_sync(
         self,
         analyze_outputs: bool = False,
+        block_run_outputs_cache: Dict[str, List] = None,
         build_block_output_stdout: Callable[..., object] = None,
         custom_code: str = None,
         execution_partition: str = None,
@@ -991,6 +999,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         output_messages_to_logs: bool = False,
         disable_json_serialization: bool = False,
         data_integration_runtime_settings: Dict = None,
+        execution_partition_previous: str = None,
         **kwargs,
     ) -> Dict:
         if logging_tags is None:
@@ -1031,6 +1040,7 @@ class Block(DataIntegrationMixin, SparkBlock):
                 )
 
             output = self.execute_block(
+                block_run_outputs_cache=block_run_outputs_cache,
                 build_block_output_stdout=build_block_output_stdout,
                 custom_code=custom_code,
                 execution_partition=execution_partition,
@@ -1044,6 +1054,7 @@ class Block(DataIntegrationMixin, SparkBlock):
                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                 run_settings=run_settings,
                 data_integration_runtime_settings=data_integration_runtime_settings,
+                execution_partition_previous=execution_partition_previous,
                 **kwargs,
             )
 
@@ -1209,6 +1220,7 @@ class Block(DataIntegrationMixin, SparkBlock):
 
     def execute_block(
         self,
+        block_run_outputs_cache: Dict[str, List] = None,
         build_block_output_stdout: Callable[..., object] = None,
         custom_code: str = None,
         execution_partition: str = None,
@@ -1223,6 +1235,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         dynamic_upstream_block_uuids: List[str] = None,
         run_settings: Dict = None,
         data_integration_runtime_settings: str = None,
+        execution_partition_previous: str = None,
         **kwargs,
     ) -> Dict:
         if logging_tags is None:
@@ -1262,11 +1275,12 @@ class Block(DataIntegrationMixin, SparkBlock):
             else:
                 input_vars, kwargs_vars, upstream_block_uuids = self.fetch_input_variables(
                     input_args,
-                    execution_partition,
-                    global_vars,
+                    block_run_outputs_cache=block_run_outputs_cache,
                     dynamic_block_index=dynamic_block_index,
                     dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                    execution_partition=execution_partition,
                     from_notebook=from_notebook,
+                    global_vars=global_vars,
                 )
 
             outputs_from_input_vars = {}
@@ -1301,6 +1315,7 @@ class Block(DataIntegrationMixin, SparkBlock):
                 upstream_block_uuids=upstream_block_uuids,
                 run_settings=run_settings,
                 data_integration_runtime_settings=data_integration_runtime_settings,
+                execution_partition_previous=execution_partition_previous,
                 **kwargs,
             )
 
@@ -1356,6 +1371,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         upstream_block_uuids: List[str] = None,
         run_settings: Dict = None,
         data_integration_runtime_settings: str = None,
+        execution_partition_previous: str = None,
         **kwargs,
     ) -> List:
         if logging_tags is None:
@@ -1379,6 +1395,7 @@ class Block(DataIntegrationMixin, SparkBlock):
                 dynamic_block_index=dynamic_block_index,
                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                 execution_partition=execution_partition,
+                execution_partition_previous=execution_partition_previous,
                 from_notebook=from_notebook,
                 global_vars=global_vars,
                 input_vars=input_vars,
@@ -1522,24 +1539,48 @@ class Block(DataIntegrationMixin, SparkBlock):
     def fetch_input_variables(
         self,
         input_args,
-        execution_partition: str = None,
-        global_vars: Dict = None,
+        block_run_outputs_cache: Dict[str, List] = None,
+        data_integration_settings_mapping: Dict = None,
         dynamic_block_index: int = None,
         dynamic_upstream_block_uuids: List[str] = None,
+        execution_partition: str = None,
         from_notebook: bool = False,
+        global_vars: Dict = None,
         upstream_block_uuids: List[str] = None,
-        data_integration_settings_mapping: Dict = None,
     ) -> Tuple[List, List, List]:
+        """
+        Fetch input variables for the current block's execution.
+
+        Args:
+            input_args: The input arguments required for the block's execution.
+            block_run_outputs_cache (Optional[Dict[str, List]]): A dictionary mapping block run
+                UUIDs to their outputs.
+            data_integration_settings_mapping (Optional[Dict]): A dictionary containing data
+                integration settings.
+            dynamic_block_index (Optional[int]): The index of the dynamic block, if applicable.
+            dynamic_upstream_block_uuids (Optional[List[str]]): The UUIDs of the dynamic upstream
+                blocks.
+            execution_partition (Optional[str]): The execution partition for the block.
+            from_notebook (Optional[bool]): A boolean indicating whether the execution is
+                triggered from a notebook.
+            global_vars (Optional[Dict]): A dictionary containing global variables.
+            upstream_block_uuids (Optional[List[str]]): List of UUIDs of upstream blocks.
+
+        Returns:
+            Tuple[List, List, List]: A tuple containing the input variables, kwargs variables, and
+                upstream block UUIDs.
+        """
         variables = fetch_input_variables(
             self.pipeline,
             upstream_block_uuids or self.upstream_block_uuids,
             input_args,
-            execution_partition,
-            global_vars,
-            dynamic_block_index,
-            dynamic_upstream_block_uuids,
-            from_notebook=from_notebook,
+            block_run_outputs_cache=block_run_outputs_cache,
             data_integration_settings_mapping=data_integration_settings_mapping,
+            dynamic_block_index=dynamic_block_index,
+            dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+            execution_partition=execution_partition,
+            from_notebook=from_notebook,
+            global_vars=global_vars,
         )
 
         return variables
@@ -2279,6 +2320,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         global_vars: Dict = None,
         logger: Logger = None,
         logging_tags: Dict = None,
+        outputs: List[Any] = None,
         update_tests: bool = True,
         dynamic_block_uuid: str = None,
     ) -> None:
@@ -2308,12 +2350,13 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         else:
             test_functions = self.test_functions
 
-        outputs = self.get_raw_outputs(
-            dynamic_block_uuid or self.uuid,
-            execution_partition=execution_partition,
-            from_notebook=from_notebook,
-            global_vars=global_vars,
-        )
+        if outputs is None:
+            outputs = self.get_raw_outputs(
+                dynamic_block_uuid or self.uuid,
+                execution_partition=execution_partition,
+                from_notebook=from_notebook,
+                global_vars=global_vars,
+            )
 
         if logger and 'logger' not in global_vars:
             global_vars['logger'] = logger
@@ -3082,8 +3125,8 @@ class ConditionalBlock(AddonBlock):
             if parent_block is not None:
                 input_vars, kwargs_vars, _ = parent_block.fetch_input_variables(
                     None,
-                    execution_partition,
-                    global_vars,
+                    execution_partition=execution_partition,
+                    global_vars=global_vars,
                     dynamic_block_index=dynamic_block_index,
                     dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                 )
@@ -3158,11 +3201,11 @@ class CallbackBlock(AddonBlock):
             # Fetch input variables
             input_vars, kwargs_vars, upstream_block_uuids = self.fetch_input_variables(
                 None,
-                execution_partition,
-                global_vars,
                 dynamic_block_index=dynamic_block_index,
                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                execution_partition=execution_partition,
                 from_notebook=from_notebook,
+                global_vars=global_vars,
                 upstream_block_uuids=[parent_block.uuid] if parent_block else None,
             )
 

mage_ai/data_preparation/models/block/data_integration/constants.py

@@ -50,6 +50,7 @@ EXECUTION_PARTITION_FROM_NOTEBOOK = '_from_notebook'
 
 OUTPUT_TYPE_RECORD = 'RECORD'
 OUTPUT_TYPE_SCHEMA = 'SCHEMA'
+OUTPUT_TYPE_STATE = 'STATE'
 TYPE_OBJECT = 'object'
 
 KEY_BOOKMARK_PROPERTIES = 'bookmark_properties'
@@ -77,6 +78,8 @@ REPLICATION_METHOD_LOG_BASED = 'LOG_BASED'
 MB_1 = 1024 * 1000
 MAX_QUERY_STRING_SIZE = 10 * MB_1
 
+VARIABLE_BOOKMARK_VALUES_KEY = '__bookmark_values__'
+
 
 class IngestMode(str, Enum):
     DISK = 'disk'

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

@@ -453,13 +453,13 @@ class DataIntegrationMixin:
         # Get the output as inputs for this block
         input_vars_fetched, kwargs_vars, up_block_uuids = self.fetch_input_variables(
             input_vars,
-            execution_partition,
-            global_vars,
+            data_integration_settings_mapping=data_integration_settings_mapping,
             dynamic_block_index=dynamic_block_index,
             dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+            execution_partition=execution_partition,
             from_notebook=from_notebook,
+            global_vars=global_vars,
             upstream_block_uuids=block_uuids_to_fetch,
-            data_integration_settings_mapping=data_integration_settings_mapping,
         )
 
         if block_uuids_to_fetch and is_debug():
@@ -515,13 +515,13 @@ class DataIntegrationMixin:
                             kwargs_vars_inner, \
                             _up_block_uuids = self.fetch_input_variables(
                                 None,
-                                execution_partition,
-                                global_vars,
+                                data_integration_settings_mapping=data_integration_settings_mapping,
                                 dynamic_block_index=dynamic_block_index,
                                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                                execution_partition=execution_partition,
                                 from_notebook=from_notebook,
+                                global_vars=global_vars,
                                 upstream_block_uuids=[up_uuid],
-                                data_integration_settings_mapping=data_integration_settings_mapping,
                             )
 
                         if input_vars_inner: