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

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:
@@ -575,11 +662,11 @@ def convert_block_output_data_for_destination(
     input_vars_fetched, _kwargs_vars, upstream_block_uuids = \
         block.fetch_input_variables(
             None,
-            partition,
-            global_vars,
             dynamic_block_index=dynamic_block_index,
             dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+            execution_partition=partition,
             from_notebook=from_notebook,
+            global_vars=global_vars,
             upstream_block_uuids=[stream],
         )
 
@@ -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/bigquery.py

@@ -29,8 +29,8 @@ def create_upstream_block_tables(
 
     input_vars, kwargs_vars, upstream_block_uuids = block.fetch_input_variables(
         None,
-        execution_partition,
-        None,
+        execution_partition=execution_partition,
+        global_vars=None,
         dynamic_block_index=dynamic_block_index,
         dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
     )

mage_ai/data_preparation/models/block/sql/snowflake.py

@@ -35,8 +35,8 @@ def create_upstream_block_tables(
 
     input_vars, kwargs_vars, upstream_block_uuids = block.fetch_input_variables(
         None,
-        execution_partition,
-        None,
+        execution_partition=execution_partition,
+        global_vars=None,
         dynamic_block_index=dynamic_block_index,
         dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
     )

mage_ai/data_preparation/models/block/sql/trino.py

@@ -32,8 +32,8 @@ def create_upstream_block_tables(
 
     input_vars, kwargs_vars, upstream_block_uuids = block.fetch_input_variables(
         None,
-        execution_partition,
-        None,
+        execution_partition=execution_partition,
+        global_vars=None,
         dynamic_block_index=dynamic_block_index,
         dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
     )

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(
@@ -353,8 +370,8 @@ def create_upstream_block_tables(
 
     input_vars, kwargs_vars, upstream_block_uuids = block.fetch_input_variables(
         None,
-        execution_partition,
-        None,
+        execution_partition=execution_partition,
+        global_vars=None,
         dynamic_block_index=dynamic_block_index,
         dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
     )

mage_ai/data_preparation/models/block/utils.py

@@ -16,6 +16,7 @@ from mage_ai.data_preparation.models.block.dynamic import (
 from mage_ai.data_preparation.models.constants import (
     DATAFRAME_ANALYSIS_MAX_COLUMNS,
     BlockType,
+    PipelineType,
 )
 from mage_ai.server.kernel_output_parser import DataType
 from mage_ai.shared.array import find, unique_by
@@ -200,6 +201,14 @@ def create_block_runs_from_dynamic_block(
                 else:
                     arr.append(upstream_block.uuid)
 
+            if metadata.get('upstream_blocks'):
+                for up_uuid in (metadata.get('upstream_blocks') or []):
+                    up_block = block.pipeline.get_block(up_uuid)
+                    if up_block:
+                        arr.append(up_uuid)
+                    else:
+                        arr.append(f'{downstream_block.uuid}:{up_uuid}')
+
             block_run = create_block_run_from_dynamic_child(
                 downstream_block,
                 pipeline_run,
@@ -513,22 +522,28 @@ def output_variables(
             partition=execution_partition,
         )
 
+    if should_reduce_output(block):
+        all_variables = all_variable_uuids(
+            block,
+            partition=execution_partition,
+        )
+    else:
+        all_variables = block.get_variables_by_block(
+            block_uuid=block_uuid,
+            partition=execution_partition,
+        )
+    output_variables = [v for v in all_variables
+                        if is_output_variable(v, include_df=include_df)]
+
     if block and di_settings:
         streams = get_selected_streams(di_settings.get('catalog'))
-        output_variables = [s.get('tap_stream_id') for s in streams]
-    else:
-        if should_reduce_output(block):
-            all_variables = all_variable_uuids(
-                block,
-                partition=execution_partition,
-            )
+        tap_stream_ids = [s.get('tap_stream_id') for s in streams]
+        # For integration pipelines, we want to include variables with the "output" prefix
+        # to view sample data for the selected streams.
+        if pipeline.type == PipelineType.INTEGRATION:
+            output_variables.extend(tap_stream_ids)
         else:
-            all_variables = block.get_variables_by_block(
-                block_uuid=block_uuid,
-                partition=execution_partition,
-            )
-        output_variables = [v for v in all_variables
-                            if is_output_variable(v, include_df=include_df)]
+            output_variables = tap_stream_ids
 
     output_variables.sort()
 
@@ -632,12 +647,13 @@ def fetch_input_variables(
     pipeline,
     upstream_block_uuids: List[str],
     input_args: List[Any],
-    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,
-    data_integration_settings_mapping: Dict = None,
+    global_vars: Dict = None,
 ) -> Tuple[List, List, List]:
     """
     Fetches the input variables for a block.
@@ -668,6 +684,7 @@ def fetch_input_variables(
             upstream_block_uuids_final = upstream_block_uuids
     elif pipeline is not None:
         input_vars = [None for i in range(len(upstream_block_uuids))]
+        # A mapping from upstream block UUID to a list of variable names
         input_variables_by_uuid = input_variables(
             pipeline,
             upstream_block_uuids,
@@ -680,6 +697,7 @@ def fetch_input_variables(
             input_args=input_args,
             data_integration_settings_mapping=data_integration_settings_mapping,
         )
+        # Block UUIDs
         keys = input_variables_by_uuid.keys()
         reduce_output_indexes = []
 
@@ -693,8 +711,10 @@ def fetch_input_variables(
                 input_vars[idx] = global_data_product.get_outputs()
                 continue
 
+            # Block output variables for upstream_block_uuid
             variables = input_variables_by_uuid[upstream_block_uuid]
 
+            # Fetch variable values
             if should_reduce:
                 variable_values = [reduce_output_from_block(
                     upstream_block,
@@ -707,19 +727,23 @@ def fetch_input_variables(
                     spark=spark,
                 ) for variable_uuid in variables]
             else:
-                variable_values = [
-                    pipeline.get_block_variable(
-                        upstream_block_uuid,
-                        var,
-                        from_notebook=from_notebook,
-                        global_vars=global_vars,
-                        input_args=input_args,
-                        partition=execution_partition,
-                        raise_exception=True,
-                        spark=spark,
-                    )
-                    for var in variables
-                ]
+                # Getting input variables from cache the cache is not empty
+                if block_run_outputs_cache:
+                    variable_values = block_run_outputs_cache.get(upstream_block_uuid, [])
+                else:
+                    variable_values = [
+                        pipeline.get_block_variable(
+                            upstream_block_uuid,
+                            var,
+                            from_notebook=from_notebook,
+                            global_vars=global_vars,
+                            input_args=input_args,
+                            partition=execution_partition,
+                            raise_exception=True,
+                            spark=spark,
+                        )
+                        for var in variables
+                    ]
 
             upstream_in_dynamic_upstream = False
             if dynamic_upstream_block_uuids:

mage_ai/data_preparation/models/global_hooks/constants.py

@@ -39,6 +39,55 @@ class HookOutputKey(str, Enum):
 
 
 VALID_KEYS_FOR_INPUT_OUTPUT_DATA_RESTRICTED = [key.value for key in HookOutputKey]
-VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED = ['hook']
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED = [
+    'hook',
+    'resource_id',
+    'resource_parent_id',
+    'user',
+]
 VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL = \
     VALID_KEYS_FOR_INPUT_OUTPUT_DATA_RESTRICTED + VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED
+
+INTERNAL_DEFAULT_PREDICATE_VALUE = '__INTERNAL_DEFAULT_PREDICATE_VALUE__'
+
+
+class PredicateAndOrOperator(str, Enum):
+    AND = 'and'
+    OR = 'or'
+
+
+class PredicateObjectType(str, Enum):
+    ERROR = HookOutputKey.ERROR.value
+    HOOK = 'hook'
+    META = HookOutputKey.META.value
+    METADATA = HookOutputKey.METADATA.value
+    OPERATION_RESOURCE = 'operation_resource'
+    PAYLOAD = HookOutputKey.PAYLOAD.value
+    QUERY = HookOutputKey.QUERY.value
+    RESOURCE = HookOutputKey.RESOURCE.value
+    RESOURCES = HookOutputKey.RESOURCES.value
+    RESOURCE_ID = 'resource_id'
+    RESOURCE_PARENT_ID = 'resource_parent_id'
+    USER = 'user'
+
+
+class PredicateOperator(str, Enum):
+    EQUALS = 'EQUALS'
+    GREATER_THAN = 'GREATER_THAN'
+    GREATER_THAN_OR_EQUALS = 'GREATER_THAN_OR_EQUALS'
+    INCLUDES = 'INCLUDES'
+    LESS_THAN = 'LESS_THAN'
+    LESS_THAN_OR_EQUALS = 'LESS_THAN_OR_EQUALS'
+    NOT_EQUALS = 'NOT_EQUALS'
+    NOT_INCLUDES = 'NOT_INCLUDES'
+    NOT_PRESENT = 'NOT_PRESENT'
+    PRESENT = 'PRESENT'
+
+
+class PredicateValueDataType(str, Enum):
+    BOOLEAN = 'BOOLEAN'
+    DICTIONARY = 'DICTIONARY'
+    FLOAT = 'FLOAT'
+    INTEGER = 'INTEGER'
+    LIST = 'LIST'
+    STRING = 'STRING'