mage-ai 0.9.45__py3-none-any.whl → 0.9.47__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

@@ -735,6 +735,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 +754,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 +777,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:

mage_ai/data_preparation/executors/pipeline_executor.py

@@ -1,6 +1,9 @@
 import asyncio
+from datetime import datetime
 from typing import Dict, List
 
+import pytz
+
 from mage_ai.data_preparation.executors.block_executor import BlockExecutor
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.logging.logger_manager_factory import LoggerManagerFactory
@@ -93,6 +96,12 @@ class PipelineExecutor:
                     block_uuid=block_run.block_uuid,
                     execution_partition=self.execution_partition,
                 )
+                block_run_data = dict(status=BlockRun.BlockRunStatus.RUNNING)
+                if not block_run.started_at or \
+                        (block_run.metrics and not block_run.metrics.get('controller')):
+                    block_run_data['started_at'] = datetime.now(tz=pytz.UTC)
+
+                block_run.update(**block_run_data)
                 BlockExecutor(**executor_kwargs).execute(
                     block_run_id=block_run.id,
                     global_vars=global_vars,

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 (
@@ -189,6 +190,7 @@ def run_blocks_sync(
     run_sensors: bool = True,
     run_tests: bool = True,
     selected_blocks: Set[str] = None,
+    update_status: bool = True,
 ) -> None:
     tries_by_block_uuid = {}
     tasks = dict()
@@ -243,6 +245,7 @@ def run_blocks_sync(
                 from_notebook=not run_sensors,
                 global_vars=global_vars,
                 run_all_blocks=True,
+                update_status=update_status,
             )
 
             if run_tests:
@@ -278,6 +281,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         language: BlockLanguage = BlockLanguage.PYTHON,
         configuration: Dict = None,
         has_callback: bool = False,
+        repo_config=None,
         timeout: int = None,
     ) -> None:
         if configuration is None:
@@ -339,12 +343,26 @@ class Block(DataIntegrationMixin, SparkBlock):
 
         self.execution_timestamp_start = None
         self.execution_timestamp_end = None
+        self.execution_uuid = None
+        self._compute_service_uuid = None
+        self._repo_config = repo_config
         self._spark_session_current = None
+        self.global_vars = None
 
     @property
     def uuid(self) -> str:
         return self._uuid
 
+    @property
+    def repo_config(self):
+        if self._repo_config:
+            return self._repo_config
+
+        if self.pipeline:
+            self._repo_config = self.pipeline.repo_config
+
+        return self._repo_config
+
     @property
     def uuid_replicated(self) -> str:
         if self.replicated_block:
@@ -449,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:
@@ -465,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():
@@ -874,6 +898,7 @@ class Block(DataIntegrationMixin, SparkBlock):
         global_vars: Dict = None,
         logger: Logger = None,
         logging_tags: Dict = None,
+        execution_uuid: str = None,
         from_notebook: bool = False,
         **kwargs
     ) -> Dict:
@@ -883,6 +908,9 @@ class Block(DataIntegrationMixin, SparkBlock):
         websocket as a way to test the code in the callback. To run a block in a pipeline
         run, use a BlockExecutor.
         """
+        if execution_uuid:
+            self.execution_uuid = execution_uuid
+
         if logging_tags is None:
             logging_tags = dict()
 
@@ -970,6 +998,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:
@@ -1023,6 +1052,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,
             )
 
@@ -1202,6 +1232,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:
@@ -1280,6 +1311,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,
             )
 
@@ -1335,6 +1367,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:
@@ -1358,6 +1391,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,
@@ -1406,8 +1440,7 @@ class Block(DataIntegrationMixin, SparkBlock):
             if logger and 'logger' not in global_vars:
                 global_vars['logger'] = logger
 
-            track_spark = from_notebook and self.is_using_spark() and \
-                self.compute_management_enabled()
+            track_spark = from_notebook and self.should_track_spark()
 
             if track_spark:
                 self.clear_spark_jobs_cache()
@@ -1490,7 +1523,8 @@ class Block(DataIntegrationMixin, SparkBlock):
             self.module = module
 
             return block_function_updated
-        except Exception:
+        except Exception as err:
+            print(f'[WARNING] Block.initialize_decorator_modules: {err}')
             print('Falling back to default block execution...')
 
         return block_function
@@ -2475,6 +2509,31 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         return variable_mapping
 
     def __enrich_global_vars(self, global_vars: Dict = None) -> Dict:
+        """
+        Enriches the provided global variables dictionary with additional context, Spark session,
+        environment, configuration, and an empty context dictionary.
+
+        Args:
+            global_vars (Optional[Dict]): A dictionary of global variables to be enriched.
+                If not provided, an empty dictionary is created.
+
+        Returns:
+            Dict: The enriched global variables dictionary.
+
+        This method checks if the pipeline type is DATABRICKS or if the environment is a Spark
+        environment. If true, it adds the Spark session to the global variables if not already
+        present.
+
+        If 'env' is not in the global variables, it adds the environment information using the
+        'get_env()' function.
+
+        Adds the block configuration to the global variables.
+
+        If 'context' is not in global_vars, it adds an empty context dictionary.
+
+        The final enriched global variables dictionary is assigned to the object's 'global_vars'
+        attribute and returned.
+        """
         if global_vars is None:
             global_vars = dict()
         if ((self.pipeline is not None and self.pipeline.type == PipelineType.DATABRICKS) or
@@ -2485,16 +2544,19 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
                     global_vars['spark'] = spark
         if 'env' not in global_vars:
             global_vars['env'] = get_env()
-        if 'configuration' not in global_vars:
-            global_vars['configuration'] = self.configuration
+        global_vars['configuration'] = self.configuration
         if 'context' not in global_vars:
             global_vars['context'] = dict()
+
+        self.global_vars = global_vars
+
         return global_vars
 
     def get_spark_session(self):
         if self.spark_init and (not self.pipeline or
                                 not self.pipeline.spark_config):
             return self.spark
+
         try:
             if self.pipeline and self.pipeline.spark_config:
                 spark_config = SparkConfig.load(
@@ -2507,6 +2569,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         except Exception:
             self.spark = None
 
+        if not self.spark and self.global_vars and self.global_vars.get('spark'):
+            self.spark = self.global_vars.get('spark')
+
         self.spark_init = True
         return self.spark
 

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/data.py

@@ -75,7 +75,7 @@ def convert_dataframe_to_output(
             if log_message:
                 log_message(
                     f'Writing {len(records_in_batch)} records from stream {stream} to {file_path} '
-                    f'for batch index {index} ({index + 1 } out of {batches} batch(s)).',
+                    f'for batch index {index} ({index + 1} out of {batches} batch(s)).',
                 )
 
             output_file_paths.append(file_path)