mage-ai 0.9.45__py3-none-any.whl → 0.9.47__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:
@@ -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/spark/mixins.py

@@ -12,6 +12,7 @@ from mage_ai.data_preparation.models.block.spark.constants import (
 from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.services.spark.api.service import API
+from mage_ai.services.spark.constants import ComputeServiceUUID
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.jobs import Job
 from mage_ai.services.spark.models.sqls import Sql
@@ -30,6 +31,16 @@ class SparkBlock:
 
         return self._spark_session_current
 
+    @property
+    def compute_service_uuid(self) -> ComputeServiceUUID:
+        if self._compute_service_uuid:
+            return self._compute_service_uuid
+        self._compute_service_uuid = get_compute_service(
+            ignore_active_kernel=True,
+            repo_config=self.repo_config,
+        )
+        return self._compute_service_uuid
+
     def spark_session_application(self) -> Application:
         if not self.spark_session:
             return
@@ -40,7 +51,7 @@ class SparkBlock:
         if value_tup:
             application_id = value_tup[1]
 
-            return Application(
+            return Application.load(
                 id=application_id,
                 spark_ui_url=self.spark_session.sparkContext.uiWebUrl,
             )
@@ -51,7 +62,13 @@ class SparkBlock:
         )
 
     def is_using_spark(self) -> bool:
-        return get_compute_service()
+        return self.compute_service_uuid in [
+            ComputeServiceUUID.AWS_EMR,
+            ComputeServiceUUID.STANDALONE_CLUSTER,
+        ]
+
+    def should_track_spark(self) -> bool:
+        return self.is_using_spark() and self.compute_management_enabled()
 
     def execution_states(self, cache: bool = False) -> Dict:
         jobs_cache = self.__load_cache()
@@ -82,33 +99,40 @@ class SparkBlock:
     def jobs_during_execution(self) -> List[Job]:
         self.__load_spark_job_submission_timestamps()
 
-        if self.execution_timestamp_start:
-            jobs = self.__get_jobs(application=self.execution_start_application)
-
-            def _filter(
-                job: Job,
-                execution_timestamp_start: float = self.execution_timestamp_start,
-                execution_timestamp_end: float = self.execution_timestamp_end,
-            ) -> bool:
-                if not job.submission_time:
-                    return False
-
-                if isinstance(job.submission_time, str):
-                    submission_timestamp = dateutil.parser.parse(job.submission_time).timestamp()
-                elif isinstance(job.submission_time, float) or isinstance(job.submission_time, int):
-                    submission_timestamp = datetime.fromtimestamp(job.submission_time)
-
-                return execution_timestamp_start and \
-                    execution_timestamp_end and \
-                    submission_timestamp >= execution_timestamp_start and \
-                    (
-                        not execution_timestamp_end or
-                        submission_timestamp <= execution_timestamp_end
-                    )
+        if not self.execution_timestamp_start and self.execution_uuid_start:
+            return []
+
+        def _filter(
+            job: Job,
+            block_uuid=self.uuid,
+            compute_service_uuid=self.compute_service_uuid,
+            execution_timestamp_end: float = self.execution_timestamp_end,
+            execution_timestamp_start: float = self.execution_timestamp_start,
+            execution_uuid_start=self.execution_uuid_start,
+        ) -> bool:
+            if ComputeServiceUUID.AWS_EMR == compute_service_uuid:
+                key = self.execution_uuid_start or self.execution_timestamp_start
+                return job.name == f'{block_uuid}:{key}'
+
+            if not job.submission_time:
+                return False
+
+            if isinstance(job.submission_time, str):
+                submission_timestamp = dateutil.parser.parse(job.submission_time).timestamp()
+            elif isinstance(job.submission_time, float) or isinstance(job.submission_time, int):
+                submission_timestamp = datetime.fromtimestamp(job.submission_time)
+
+            return execution_timestamp_start and \
+                execution_timestamp_end and \
+                submission_timestamp >= execution_timestamp_start and \
+                (
+                    not execution_timestamp_end or
+                    submission_timestamp <= execution_timestamp_end
+                )
 
-            return list(filter(_filter, jobs))
+        jobs = self.__get_jobs(application=self.execution_start_application)
 
-        return []
+        return list(filter(_filter, jobs or []))
 
     def stages_during_execution(self, jobs: List[Job]):
         if not jobs:
@@ -122,7 +146,7 @@ class SparkBlock:
 
         stages = self.__get_stages(application=self.execution_start_application)
 
-        return list(filter(_filter, stages))
+        return list(filter(_filter, stages or []))
 
     def sqls_during_execution(self, jobs: List[Job]):
         if not jobs:
@@ -140,7 +164,7 @@ class SparkBlock:
 
         sqls = self.__get_sqls(application=self.execution_start_application)
 
-        return list(filter(_filter, sqls))
+        return list(filter(_filter, sqls or []))
 
     def clear_spark_jobs_cache(self) -> None:
         if os.path.exists(self.spark_jobs_full_path):
@@ -155,13 +179,24 @@ class SparkBlock:
         for application in applications:
             Application.cache_application(application)
 
-    def set_spark_job_execution_start(self) -> None:
+    def set_spark_job_execution_start(self, execution_uuid: str = None) -> None:
         self.execution_timestamp_start = datetime.utcnow().timestamp()
         application = self.spark_session_application()
 
+        if execution_uuid:
+            self.execution_uuid = execution_uuid
+
+        if self.spark_session and self.spark_session.sparkContext:
+            key = f'{self.uuid}:{self.execution_uuid or self.execution_timestamp_start}'
+            # For jobs
+            self.spark_session.sparkContext.setLocalProperty('callSite.short', key)
+            # For stages
+            self.spark_session.sparkContext.setLocalProperty('callSite.long', key)
+
         self.__update_spark_jobs_cache(
             dict(
                 application=application.to_dict() if application else None,
+                execution_uuid=self.execution_uuid,
                 submission_timestamp=self.execution_timestamp_start,
             ),
             'before',
@@ -187,7 +222,7 @@ class SparkBlock:
 
         if application:
             build_options.update(dict(
-                application_id=application.id,
+                application_id=application.calculated_id(),
                 application_spark_ui_url=application.spark_ui_url,
             ))
         else:
@@ -207,7 +242,7 @@ class SparkBlock:
         else:
             applications = api.applications_sync()
             if applications:
-                jobs = api.jobs_sync(applications[0].id)
+                jobs = api.jobs_sync(applications[0].calculated_id())
 
         return sorted(
             jobs,
@@ -231,7 +266,7 @@ class SparkBlock:
             applications = api.applications_sync()
 
             if applications:
-                stages = api.stages_sync(applications[0].id, dict(
+                stages = api.stages_sync(applications[0].calculated_id(), dict(
                     quantiles='0.01,0.25,0.5,0.75,0.99',
                     withSummaries=True,
                 ))
@@ -253,7 +288,7 @@ class SparkBlock:
 
             sqls = []
             if applications:
-                sqls = api.sqls_sync(applications[0].id, dict(
+                sqls = api.sqls_sync(applications[0].calculated_id(), dict(
                     length=9999,
                 ))
 
@@ -304,6 +339,13 @@ class SparkBlock:
             f.write(json.dumps(data))
 
     def __load_spark_job_submission_timestamps(self) -> None:
+        self.execution_end_application = None
+        self.execution_start_application = None
+        self.execution_timestamp_end = None
+        self.execution_timestamp_start = None
+        self.execution_uuid_end = None
+        self.execution_uuid_start = None
+
         jobs_cache = self.__load_cache()
         if jobs_cache:
             before = jobs_cache.get('before')
@@ -311,6 +353,9 @@ class SparkBlock:
                 self.execution_timestamp_start = (before or {}).get(
                     'submission_timestamp',
                 )
+                self.execution_uuid_start = (before or {}).get(
+                    'execution_uuid',
+                )
                 self.execution_start_application = (before or {}).get(
                     'application',
                 )
@@ -324,6 +369,9 @@ class SparkBlock:
                 self.execution_timestamp_end = (after or {}).get(
                     'submission_timestamp',
                 )
+                self.execution_uuid_end = (after or {}).get(
+                    'execution_uuid',
+                )
                 self.execution_end_application = (after or {}).get(
                     'application',
                 )

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(

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()
 

mage_ai/data_preparation/models/download/__init__.py

@@ -0,0 +1,8 @@
+class Download():
+    def __init__(self, token):
+        self.token = token
+
+    def to_dict(self, **kwargs):
+        data = {}
+        data['token'] = self.token
+        return data