mage-ai 0.9.46__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/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():
@@ -991,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:
@@ -1044,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,
             )
 
@@ -1223,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:
@@ -1301,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,
             )
 
@@ -1356,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:
@@ -1379,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,

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

@@ -24,6 +24,7 @@ from mage_ai.data_preparation.models.block.data_integration.constants import (
     KEY_METADATA,
     KEY_PARTITION_KEYS,
     KEY_PROPERTIES,
+    KEY_RECORD,
     KEY_REPLICATION_METHOD,
     KEY_SCHEMA,
     KEY_STREAM,
@@ -31,10 +32,14 @@ from mage_ai.data_preparation.models.block.data_integration.constants import (
     KEY_TYPE,
     KEY_UNIQUE_CONFLICT_METHOD,
     KEY_UNIQUE_CONSTRAINTS,
+    KEY_VALUE,
     MAX_QUERY_STRING_SIZE,
+    OUTPUT_TYPE_RECORD,
     OUTPUT_TYPE_SCHEMA,
+    OUTPUT_TYPE_STATE,
     REPLICATION_METHOD_INCREMENTAL,
     STATE_FILENAME,
+    VARIABLE_BOOKMARK_VALUES_KEY,
     IngestMode,
 )
 from mage_ai.data_preparation.models.block.data_integration.data import (
@@ -48,6 +53,7 @@ from mage_ai.data_preparation.models.constants import (
 )
 from mage_ai.data_preparation.models.pipelines.utils import number_string
 from mage_ai.shared.array import find
+from mage_ai.shared.files import reverse_readline
 from mage_ai.shared.hash import dig, extract, merge_dict
 from mage_ai.shared.parsers import encode_complex, extract_json_objects
 from mage_ai.shared.security import filter_out_config_values
@@ -278,23 +284,82 @@ def get_streams_from_output_directory(
             # or
             # ../[block_uuid]/[data_integration_uuid]
             dir_full_path1 = os.path.join(output_directory_path, dir_name1)
-            for dir_name2 in os.listdir(dir_full_path1):
-                if not data_integration_uuid:
-                    mapping[dir_name2] = []
-                # ../[block_uuid]/[data_integration_uuid]/[stream]/00000000000000000000
-                # or
-                # ../[block_uuid]/[data_integration_uuid]/[stream]
-                dir_full_path2 = os.path.join(dir_full_path1, dir_name2)
-                if data_integration_uuid:
-                    mapping[dir_name1].append(dir_full_path2)
-                else:
-                    for dir_name3 in os.listdir(dir_full_path2):
-                        dir_full_path3 = os.path.join(dir_full_path2, dir_name3)
-                        mapping[dir_name2].append(dir_full_path3)
+            if os.path.isdir(dir_full_path1):
+                for dir_name2 in os.listdir(dir_full_path1):
+                    if not data_integration_uuid:
+                        mapping[dir_name2] = []
+                    # ../[block_uuid]/[data_integration_uuid]/[stream]/00000000000000000000
+                    # or
+                    # ../[block_uuid]/[data_integration_uuid]/[stream]
+                    dir_full_path2 = os.path.join(dir_full_path1, dir_name2)
+                    if data_integration_uuid:
+                        mapping[dir_name1].append(dir_full_path2)
+                    elif os.path.isdir(dir_full_path2):
+                        for dir_name3 in os.listdir(dir_full_path2):
+                            dir_full_path3 = os.path.join(dir_full_path2, dir_name3)
+                            mapping[dir_name2].append(dir_full_path3)
 
     return mapping
 
 
+def get_state_data(
+    block,
+    catalog: Dict,
+    from_notebook: bool = False,
+    index: int = None,
+    partition: str = None,
+    data_integration_uuid: str = None,
+    include_record: bool = False,
+    stream_id: str = None,
+) -> Union[Dict, Tuple[Dict, Dict]]:
+    output_file_paths = get_output_file_paths(
+        block,
+        catalog,
+        from_notebook=from_notebook,
+        index=index,
+        partition=partition,
+        data_integration_uuid=data_integration_uuid,
+        stream_id=stream_id,
+    )
+
+    output_file_paths.sort()
+
+    record = None
+    state_data = None
+
+    if output_file_paths:
+        output_file_path = output_file_paths[-1]
+
+        for line in reverse_readline(output_file_path):
+            if line:
+                try:
+                    row = json.loads(line)
+                    row_type = row.get(KEY_TYPE)
+
+                    if include_record and \
+                            OUTPUT_TYPE_RECORD == row_type and \
+                            KEY_RECORD in row and \
+                            (not stream_id or stream_id == row.get(KEY_STREAM)):
+
+                        record = row[KEY_RECORD]
+                    elif OUTPUT_TYPE_STATE == row_type and KEY_VALUE in row:
+                        # If it finds a state again even before it find a record, break.
+                        if state_data is not None:
+                            break
+
+                        state_data = row[KEY_VALUE]
+
+                        if not include_record or record:
+                            break
+                except json.JSONDecodeError:
+                    pass
+
+        if include_record:
+            return state_data, record
+
+        return state_data
+
+
 def execute_data_integration(
     block,
     outputs_from_input_vars,
@@ -303,6 +368,7 @@ def execute_data_integration(
     dynamic_block_index: int = None,
     dynamic_upstream_block_uuids: List[str] = None,
     execution_partition: str = None,
+    execution_partition_previous: str = None,
     from_notebook: bool = False,
     global_vars: Dict = None,
     input_from_output: Dict = None,
@@ -382,17 +448,34 @@ def execute_data_integration(
             return []
 
         # Handle incremental sync
-        state_file_path = None
+        state_data = None
         if index is not None:
             batch_fetch_limit = get_batch_fetch_limit(config)
-            state_file_path = get_state_file_path(block, data_integration_uuid, stream)
             stream_catalogs = get_streams_from_catalog(catalog, [stream]) or []
 
             if len(stream_catalogs) == 1 and \
                     REPLICATION_METHOD_INCREMENTAL == stream_catalogs[0].get('replication_method'):
-                # Use the state to adjust the query
-                # How do we write to the state when the source syncs can run in parallel?
-                pass
+
+                if global_vars_more and VARIABLE_BOOKMARK_VALUES_KEY in global_vars_more:
+                    bookmark_values_by_block_uuid = global_vars_more.get(
+                        VARIABLE_BOOKMARK_VALUES_KEY,
+                    ) or {}
+
+                    if bookmark_values_by_block_uuid.get(block.uuid):
+                        state_data = dict(
+                            bookmarks=bookmark_values_by_block_uuid.get(block.uuid),
+                        )
+
+                if not state_data and execution_partition_previous:
+                    state_data = get_state_data(
+                        block,
+                        catalog,
+                        data_integration_uuid=data_integration_uuid,
+                        from_notebook=from_notebook,
+                        index=index,
+                        partition=execution_partition_previous,
+                        stream_id=stream,
+                    )
             else:
                 query_data['_offset'] = batch_fetch_limit * index
 
@@ -436,10 +519,14 @@ def execute_data_integration(
                 block.get_catalog_file_path(),
             ]
 
-        if state_file_path:
+        if state_data:
             args += [
-                '--state',
-                state_file_path,
+                '--state_json',
+                simplejson.dumps(
+                    state_data,
+                    default=encode_complex,
+                    ignore_nan=True,
+                ),
             ]
 
         if len(selected_streams) >= 1:
@@ -953,7 +1040,7 @@ def __execute_destination(
     return proc
 
 
-def convert_outputs_to_data(
+def get_output_file_paths(
     block,
     catalog: Dict,
     from_notebook: bool = False,
@@ -961,8 +1048,7 @@ def convert_outputs_to_data(
     partition: str = None,
     data_integration_uuid: str = None,
     stream_id: str = None,
-    sample_count: int = None,
-) -> Dict:
+) -> List[str]:
     variable = build_variable(
         block,
         data_integration_uuid=data_integration_uuid,
@@ -987,6 +1073,29 @@ def convert_outputs_to_data(
         output_file_paths.append(output_file_path)
     output_file_paths.sort()
 
+    return output_file_paths
+
+
+def convert_outputs_to_data(
+    block,
+    catalog: Dict,
+    from_notebook: bool = False,
+    index: int = None,
+    partition: str = None,
+    data_integration_uuid: str = None,
+    stream_id: str = None,
+    sample_count: int = None,
+) -> Dict:
+    output_file_paths = get_output_file_paths(
+        block,
+        catalog,
+        from_notebook=from_notebook,
+        index=index,
+        partition=partition,
+        data_integration_uuid=data_integration_uuid,
+        stream_id=stream_id,
+    )
+
     columns_to_select = []
     rows = []
     stream_settings = {}
@@ -1011,19 +1120,20 @@ def convert_outputs_to_data(
         columns_to_select = [d.get('column') for d in columns]
 
         for output_file_path in output_file_paths:
-            with open(output_file_path) as f:
-                for line in f:
-                    try:
-                        if sample_count is not None and row_count >= sample_count:
-                            break
-
-                        row = json.loads(line)
-                        record = row.get('record')
-                        if record and stream_id == row.get('stream'):
-                            rows.append([record.get(col) for col in columns_to_select])
-                            row_count += 1
-                    except json.JSONDecodeError:
-                        pass
+            if os.path.exists(output_file_path):
+                with open(output_file_path) as f:
+                    for line in f:
+                        try:
+                            if sample_count is not None and row_count >= sample_count:
+                                break
+
+                            row = json.loads(line)
+                            record = row.get('record')
+                            if record and stream_id == row.get('stream'):
+                                rows.append([record.get(col) for col in columns_to_select])
+                                row_count += 1
+                        except json.JSONDecodeError:
+                            pass
 
     if sample_count is not None:
         rows = rows[:sample_count]
@@ -1113,8 +1223,12 @@ def count_records(
     config: Dict,
     source_uuid: str,
     streams: List[str],
+    block=None,
     catalog: Dict = None,
     catalog_file_path: str = None,
+    from_notebook: bool = False,
+    partition: str = None,
+    variables: Dict = None,
 ) -> List[Dict]:
     arr = []
 
@@ -1152,6 +1266,51 @@ def count_records(
                 catalog_file_path,
             ]
 
+        state_data = None
+        # Make sure replication method is INCREMENTAL
+        if block:
+            stream_dicts = catalog.get('streams')
+            if stream_dicts:
+                stream_dict = find(lambda x, stream=stream: x['stream'] == stream, stream_dicts)
+                if not stream_dict:
+                    raise Exception(
+                        f'No stream settings found for stream {stream} in source {source_uuid}, '
+                        'this is unexpected.',
+                    )
+
+                if stream_dict and \
+                        REPLICATION_METHOD_INCREMENTAL == stream_dict.get(KEY_REPLICATION_METHOD):
+
+                    if variables and VARIABLE_BOOKMARK_VALUES_KEY in variables:
+                        bookmark_values_by_block_uuid = variables.get(
+                            VARIABLE_BOOKMARK_VALUES_KEY,
+                        ) or {}
+
+                        if bookmark_values_by_block_uuid.get(block.uuid):
+                            state_data = dict(
+                                bookmarks=bookmark_values_by_block_uuid.get(block.uuid),
+                            )
+
+                    if not state_data:
+                        state_data = get_state_data(
+                            block,
+                            catalog,
+                            data_integration_uuid=source_uuid,
+                            from_notebook=from_notebook,
+                            partition=partition,
+                            stream_id=stream,
+                        )
+
+        if state_data:
+            args += [
+                '--state_json',
+                simplejson.dumps(
+                    state_data,
+                    default=encode_complex,
+                    ignore_nan=True,
+                ),
+            ]
+
         arr += json.loads(__run_in_subprocess(args, config=config))
 
     return arr

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

@@ -10,9 +10,12 @@ from mage_ai.data_preparation.models.block.sql.constants import (
     CONFIG_KEY_UPSTREAM_BLOCK_CONFIGURATION_TABLE_NAME,
 )
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
+from mage_ai.data_preparation.shared.utils import get_template_vars
+from mage_ai.data_preparation.templates.utils import get_variable_for_template
 from mage_ai.data_preparation.variable_manager import get_variable
 from mage_ai.io.config import ConfigFileLoader
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.hash import merge_dict
 
 MAGE_SEMI_COLON = '__MAGE_SEMI_COLON__'
 
@@ -228,7 +231,21 @@ def interpolate_input(
 def interpolate_vars(query, global_vars=None):
     if global_vars is None:
         global_vars = dict()
-    return Template(query, undefined=StrictUndefined).render(**global_vars)
+
+    return Template(
+        query,
+        undefined=StrictUndefined,
+    ).render(
+        variables=lambda x, p=None, v=global_vars: get_variable_for_template(
+            x,
+            parse=p,
+            variables=v,
+        ),
+        **merge_dict(
+            global_vars,
+            get_template_vars(),
+        ),
+    )
 
 
 def table_name_parts(