mage-ai 0.9.75__py3-none-any.whl → 0.9.77__py3-none-any.whl

mage_ai/api/policies/PipelineSchedulePolicy.py

@@ -100,6 +100,7 @@ PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes +
 PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes + [
     'event_matchers',
     'next_pipeline_run_date',
+    'rotate_token',
     'tags',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/presenters/PipelineSchedulePresenter.py

@@ -1,5 +1,6 @@
 from mage_ai.api.operations import constants
 from mage_ai.api.presenters.BasePresenter import BasePresenter
+from mage_ai.settings.server import HIDE_API_TRIGGER_TOKEN
 
 
 class PipelineSchedulePresenter(BasePresenter):
@@ -32,8 +33,6 @@ class PipelineSchedulePresenter(BasePresenter):
             ])
             data['tags'] = sorted([tag.name for tag in self.get_tag_associations])
             data['next_pipeline_run_date'] = self.model.next_execution_date()
-
-            return data
         elif 'with_runtime_average' == display_format:
             data = self.model.to_dict()
             data['runtime_average'] = self.model.runtime_average()
@@ -47,6 +46,16 @@ class PipelineSchedulePresenter(BasePresenter):
         else:
             data = self.model.to_dict()
 
+        if display_format == constants.UPDATE:
+            rotate_token = kwargs.get(
+                'payload', dict(),
+            ).get(
+                'pipeline_schedule', dict(),
+            ).get('rotate_token')
+        else:
+            rotate_token = False
+        if HIDE_API_TRIGGER_TOKEN and not rotate_token:
+            data['token'] = '[API_TOKEN_PLACEHOLDER]'
         return data
 
 

mage_ai/api/resources/GitFileResource.py

@@ -32,6 +32,14 @@ class GitFileResource(GenericResource):
                 pass
 
         file_path_absolute = os.path.join(git_manager.repo_path, file_path)
+
+        # Prevent path traversal by resolving the absolute path location
+        # and checking if it's within the repo
+        if not os.path.abspath(file_path_absolute).startswith(
+            os.path.abspath(git_manager.repo_path)
+        ):
+            raise Exception("Access denied: Attempted path traversal")
+
         file = File.from_path(file_path_absolute)
         if not file.exists():
             file = File.from_path(file_path_absolute, '')

mage_ai/api/resources/PipelineScheduleResource.py

@@ -29,7 +29,8 @@ from mage_ai.orchestration.db.models.tags import (
 )
 from mage_ai.settings.platform import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
-from mage_ai.shared.hash import merge_dict
+from mage_ai.settings.server import HIDE_API_TRIGGER_TOKEN
+from mage_ai.shared.hash import ignore_keys, merge_dict
 
 
 class PipelineScheduleResource(DatabaseResource):
@@ -282,6 +283,7 @@ class PipelineScheduleResource(DatabaseResource):
 
     @safe_db_query
     def update(self, payload, **kwargs):
+        # Update associated event matchers
         arr = payload.pop('event_matchers', None)
         event_matchers = []
         if arr is not None:
@@ -320,6 +322,7 @@ class PipelineScheduleResource(DatabaseResource):
                 ]
                 em.update(pipeline_schedules=ps)
 
+        # Update associated tags
         tag_names = payload.pop('tags', None)
         if tag_names is not None:
             # 1. Fetch all tag associations
@@ -408,7 +411,11 @@ class PipelineScheduleResource(DatabaseResource):
 
         old_name = self.model.name
 
-        resource = super().update(payload)
+        # Rotate token
+        if payload.get('rotate_token'):
+            payload['token'] = uuid.uuid4().hex
+
+        resource = super().update(ignore_keys(payload, ['rotate_token']))
         updated_model = resource.model
 
         repo_path = get_repo_path(user=self.current_user)
@@ -424,9 +431,10 @@ class PipelineScheduleResource(DatabaseResource):
                 sla=updated_model.sla,
                 start_time=updated_model.start_time,
                 status=updated_model.status,
-                token=updated_model.token,
                 variables=updated_model.variables,
             )
+            if not HIDE_API_TRIGGER_TOKEN:
+                trigger.token = updated_model.token
 
             update_only_if_exists = (
                 not pipeline.should_save_trigger_in_code_automatically()

mage_ai/api/resources/PipelineTriggerResource.py

@@ -8,6 +8,7 @@ from mage_ai.data_preparation.models.triggers import (
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.schedules import PipelineSchedule
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.settings.server import HIDE_API_TRIGGER_TOKEN
 
 
 class PipelineTriggerResource(GenericResource):
@@ -50,9 +51,10 @@ class PipelineTriggerResource(GenericResource):
                 sla=pipeline_schedule.sla,
                 start_time=pipeline_schedule.start_time,
                 status=pipeline_schedule.status,
-                token=pipeline_schedule.token,
                 variables=pipeline_schedule.variables,
             )
+            if not HIDE_API_TRIGGER_TOKEN:
+                trigger.token = pipeline_schedule.token
         else:
             trigger = Trigger(**payload)
 

mage_ai/api/resources/SessionResource.py

@@ -6,7 +6,7 @@ from mage_ai.authentication.ldap import new_ldap_connection
 from mage_ai.authentication.oauth2 import encode_token, generate_access_token
 from mage_ai.authentication.passwords import verify_password
 from mage_ai.authentication.providers.constants import NAME_TO_PROVIDER
-from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db import safe_db_query, safe_db_query_async
 from mage_ai.orchestration.db.models.oauth import Role, User
 from mage_ai.settings import (
     AUTHENTICATION_MODE,
@@ -20,7 +20,7 @@ from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 class SessionResource(BaseResource):
     @classmethod
-    @safe_db_query
+    @safe_db_query_async
     async def create(self, payload, _, **kwargs):
         email = payload.get('email')
         password = payload.get('password')

mage_ai/api/resources/SyncResource.py

@@ -163,7 +163,7 @@ class SyncResource(GenericResource):
         access_token = user_payload.pop('access_token', None)
         if access_token:
             secret_name = get_access_token_secret_name(user=user)
-            secret = Secret.query.filter(Secret.name == secret_name).one_or_none()
+            secret = Secret.repo_query.filter(Secret.name == secret_name).one_or_none()
             if secret:
                 secret.delete()
             create_secret(secret_name, access_token, repo_name=repo_name)

mage_ai/api/resources/UserResource.py

@@ -28,7 +28,7 @@ class UserResource(DatabaseResource):
         results = (
             User.
             query.
-            order_by(User.username.asc())
+            order_by(User.id.asc())
         )
 
         if user and user.is_admin:

mage_ai/cli/main.py

@@ -194,7 +194,11 @@ def run(
     from mage_ai.orchestration.db.models.schedules import PipelineRun
     from mage_ai.orchestration.utils.git import log_git_sync, run_git_sync
     from mage_ai.server.logger import Logger
-    from mage_ai.settings import SENTRY_DSN, SENTRY_TRACES_SAMPLE_RATE
+    from mage_ai.settings import (
+        SENTRY_DSN,
+        SENTRY_SERVER_NAME,
+        SENTRY_TRACES_SAMPLE_RATE,
+    )
     from mage_ai.shared.hash import merge_dict
 
     logger = Logger().new_server_logger(__name__)
@@ -204,6 +208,7 @@ def run(
         sentry_sdk.init(
             sentry_dsn,
             traces_sample_rate=SENTRY_TRACES_SAMPLE_RATE,
+            server_name=SENTRY_SERVER_NAME,
         )
     (enable_new_relic, application) = initialize_new_relic()
 

mage_ai/data_integrations/destinations/constants.py

@@ -5,6 +5,7 @@ DESTINATIONS = [
     dict(name='Clickhouse'),
     dict(name='Delta Lake Azure'),
     dict(name='Delta Lake S3'),
+    dict(name='Doris'),
     dict(name='Elasticsearch'),
     dict(name='Google Cloud Storage'),
     dict(name='Kafka'),
@@ -21,5 +22,6 @@ DESTINATIONS = [
     dict(name='Redshift'),
     dict(name='Salesforce'),
     dict(name='Snowflake'),
+    dict(name='Teradata'),
     dict(name='Trino'),
 ]

mage_ai/data_integrations/sources/constants.py

@@ -3,6 +3,7 @@ from mage_ai.shared.hash import index_by
 
 SQL_SOURCES = [
     dict(name='BigQuery'),
+    dict(name='Doris'),
     dict(
         name='Microsoft SQL Server',
         uuid='mssql',
@@ -53,6 +54,7 @@ SOURCES = sorted([
     dict(name='Sftp'),
     dict(name='Stripe'),
     dict(name='Tableau'),
+    dict(name='Teradata'),
     dict(name='Twitter Ads'),
     dict(name='Zendesk'),
 ] + SQL_SOURCES, key=lambda x: x['name'])

mage_ai/data_preparation/executors/block_executor.py

@@ -693,7 +693,7 @@ class BlockExecutor:
                             ),
                             tags=tags,
                         )
-                    self._execute_callback(
+                    self.execute_callback(
                         'on_failure',
                         block_run_id=block_run_id,
                         callback_kwargs=dict(__error=error, retry=self.retry_metadata),
@@ -748,7 +748,7 @@ class BlockExecutor:
             # success callback because this isn’t the last data integration block that needs
             # to run.
             if not data_integration_metadata or is_original_block:
-                self._execute_callback(
+                self.execute_callback(
                     'on_success',
                     block_run_id=block_run_id,
                     callback_kwargs=dict(retry=self.retry_metadata),
@@ -1253,7 +1253,7 @@ class BlockExecutor:
 
         return result
 
-    def _execute_callback(
+    def execute_callback(
         self,
         callback: str,
         global_vars: Dict,
@@ -1275,6 +1275,11 @@ class BlockExecutor:
             dynamic_block_index: Index of the dynamic block.
             dynamic_upstream_block_uuids: List of UUIDs of the dynamic upstream blocks.
         """
+        if logging_tags is None:
+            logging_tags = self.build_tags(
+                block_run_id=block_run_id,
+                pipeline_run_id=pipeline_run.id if pipeline_run is not None else None,
+            )
         upstream_block_uuids_override = []
         if is_dynamic_block_child(self.block):
             if not self.block_run and block_run_id:

mage_ai/data_preparation/executors/pipeline_executor.py

@@ -10,10 +10,13 @@ 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.pipeline import Pipeline
 from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun
+from mage_ai.server.logger import Logger
 from mage_ai.shared.hash import merge_dict
 from mage_ai.usage_statistics.constants import EventNameType, EventObjectType
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
+logger = Logger().new_server_logger(__name__)
+
 
 class PipelineExecutor:
     def __init__(self, pipeline: Pipeline, execution_partition: str = None):
@@ -53,25 +56,38 @@ class PipelineExecutor:
             update_status (bool): Whether to update the execution status.
             **kwargs: Additional keyword arguments.
         """
-        if pipeline_run_id is None:
-            # Execute the pipeline without block runs
-            asyncio.run(self.pipeline.execute(
-                analyze_outputs=analyze_outputs,
-                global_vars=global_vars,
-                run_sensors=run_sensors,
-                run_tests=run_tests,
-                update_status=update_status,
-            ))
-        else:
-            # Supported pipeline types: Standard batch pipeline
-            pipeline_run = PipelineRun.query.get(pipeline_run_id)
-            if pipeline_run.status != PipelineRun.PipelineRunStatus.RUNNING:
-                return
-            asyncio.run(self.__run_blocks(
-                pipeline_run,
-                allow_blocks_to_fail=allow_blocks_to_fail,
-                global_vars=global_vars,
-            ))
+        # Create the async task to execute
+        async def _execute_task():
+            if pipeline_run_id is None:
+                # Execute the pipeline without block runs
+                await self.pipeline.execute(
+                    analyze_outputs=analyze_outputs,
+                    global_vars=global_vars,
+                    run_sensors=run_sensors,
+                    run_tests=run_tests,
+                    update_status=update_status,
+                )
+            else:
+                # Supported pipeline types: Standard batch pipeline
+                pipeline_run = PipelineRun.query.get(pipeline_run_id)
+                if pipeline_run.status != PipelineRun.PipelineRunStatus.RUNNING:
+                    return
+                await self.__run_blocks(
+                    pipeline_run,
+                    allow_blocks_to_fail=allow_blocks_to_fail,
+                    global_vars=global_vars,
+                )
+        # Execute the task based on current context
+        try:
+            loop = asyncio.get_running_loop()
+            logger.info(f'[PipelineExecutor] Found running loop {loop}')
+            # We're in an async context, use create_task
+            task = asyncio.create_task(_execute_task())
+            loop.run_until_complete(task)
+        except RuntimeError:
+            # No running loop, safe to use asyncio.run
+            logger.info('[PipelineExecutor] No running loop, using asyncio.run')
+            asyncio.run(_execute_task())
 
         self.logger_manager.output_logs_to_destination()
 

mage_ai/data_preparation/models/block/__init__.py

@@ -3941,7 +3941,6 @@ class Block(
             for v in b.output_variables(execution_partition=execution_partition):
                 objs.append(
                     self.get_variable_object(
-                        self.pipeline_uuid,
                         b.uuid,
                         v,
                         partition=execution_partition,
@@ -4465,29 +4464,36 @@ class CallbackBlock(AddonBlock):
             elif 'on_success' == callback:
                 callback_functions_legacy = success_functions
                 callback_status = CallbackStatus.SUCCESS
+            elif 'on_cancelled' == callback:
+                callback_functions_legacy = []
+                callback_status = CallbackStatus.CANCELLED
 
-            # Fetch input variables
-            input_vars, kwargs_vars, upstream_block_uuids = self.fetch_input_variables(
-                None,
-                dynamic_block_index=dynamic_block_index,
-                dynamic_block_indexes=dynamic_block_indexes,
-                dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
-                execution_partition=execution_partition,
-                from_notebook=from_notebook,
-                global_vars=global_vars,
-                metadata=metadata,
-                upstream_block_uuids=[parent_block.uuid] if parent_block else None,
-                upstream_block_uuids_override=upstream_block_uuids_override,
-            )
-
-            # Copied logic from the method self.execute_block
-            outputs_from_input_vars = {}
-            upstream_block_uuids_length = len(upstream_block_uuids)
-            for idx, input_var in enumerate(input_vars):
-                if idx < upstream_block_uuids_length:
-                    upstream_block_uuid = upstream_block_uuids[idx]
-                    outputs_from_input_vars[upstream_block_uuid] = input_var
-                    outputs_from_input_vars[f'df_{idx + 1}'] = input_var
+            if callback_status in [CallbackStatus.FAILURE, CallbackStatus.SUCCESS]:
+                # Fetch input variables
+                input_vars, kwargs_vars, upstream_block_uuids = self.fetch_input_variables(
+                    None,
+                    dynamic_block_index=dynamic_block_index,
+                    dynamic_block_indexes=dynamic_block_indexes,
+                    dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                    execution_partition=execution_partition,
+                    from_notebook=from_notebook,
+                    global_vars=global_vars,
+                    metadata=metadata,
+                    upstream_block_uuids=[parent_block.uuid] if parent_block else None,
+                    upstream_block_uuids_override=upstream_block_uuids_override,
+                )
+                # Copied logic from the method self.execute_block
+                outputs_from_input_vars = {}
+                upstream_block_uuids_length = len(upstream_block_uuids)
+                for idx, input_var in enumerate(input_vars):
+                    if idx < upstream_block_uuids_length:
+                        upstream_block_uuid = upstream_block_uuids[idx]
+                        outputs_from_input_vars[upstream_block_uuid] = input_var
+                        outputs_from_input_vars[f'df_{idx + 1}'] = input_var
+            else:
+                input_vars = []
+                kwargs_vars = []
+                upstream_block_uuids = []
 
             global_vars_copy = global_vars.copy()
             for kwargs_var in kwargs_vars:

mage_ai/data_preparation/models/block/dbt/dbt_adapter.py

@@ -5,7 +5,9 @@ from typing import Any, Dict, Optional, Tuple, Union
 
 import dbt.flags as flags
 import pandas as pd
-from dbt.adapters.base import BaseRelation, Credentials
+from dbt.adapters.base.relation import BaseRelation
+from dbt.adapters.contracts.connection import AdapterResponse, Credentials
+from dbt.adapters.contracts.relation import RelationType
 from dbt.adapters.factory import (
     Adapter,
     cleanup_connections,
@@ -13,10 +15,10 @@ from dbt.adapters.factory import (
     register_adapter,
     reset_adapters,
 )
-from dbt.config.profile import read_user_config
+from dbt.config.project import read_project_flags
 from dbt.config.runtime import RuntimeConfig
-from dbt.contracts.connection import AdapterResponse
-from dbt.contracts.relation import RelationType
+from dbt.context.providers import generate_runtime_macro_context
+from dbt.mp_context import get_mp_context
 
 from mage_ai.data_preparation.models.block.dbt.profiles import Profiles
 from mage_ai.shared.environments import is_debug
@@ -79,7 +81,11 @@ class DBTAdapter:
         # remove interpolated profiles.yml
         self.__profiles.clean()
 
-    def execute(self, sql: str, fetch: bool = False) -> Tuple[AdapterResponse, pd.DataFrame]:
+    def execute(
+        self,
+        sql: str,
+        fetch: bool = False
+    ) -> Tuple[AdapterResponse, pd.DataFrame]:
         """
         Executes any sql statement using the dbt adapter.
 
@@ -135,6 +141,8 @@ class DBTAdapter:
             package
         )
 
+        self.__adapter.set_macro_resolver(manifest)
+
         # create a context for the macro (e.g. downstream macros)
         from dbt.context.providers import generate_runtime_macro_context
         macro_context = generate_runtime_macro_context(
@@ -203,7 +211,7 @@ class DBTAdapter:
         # set dbt flags
         # Need to add profiles.yml file
         try:
-            user_config = read_user_config(profiles_path)
+            user_config = read_project_flags(self.project_path, profiles_path)
         except Exception as err:
             print(f'[ERROR] DBTAdapter.open: {err}.')
 
@@ -211,7 +219,10 @@ class DBTAdapter:
                     not profiles_path.endswith('profiles.yml'):
 
                 try:
-                    user_config = read_user_config(os.path.join(profiles_path, 'profiles.yml'))
+                    user_config = read_project_flags(
+                        self.project_path,
+                        os.path.join(profiles_path, 'profiles.yml')
+                    )
                 except Exception as err2:
                     print(f'[ERROR] DBTAdapter.open: {err2}.')
                     raise err
@@ -227,9 +238,10 @@ class DBTAdapter:
             config = RuntimeConfig.from_args(adapter_config)
             reset_adapters()
             # register the correct adapter from config
-            register_adapter(config)
+            register_adapter(config, mp_context=get_mp_context())
             # load the adapter
             self.__adapter = get_adapter(config)
+            self.__adapter.set_macro_context_generator(generate_runtime_macro_context)
             # connect
             self.__adapter.acquire_connection('mage_dbt_adapter_' + uuid.uuid4().hex)
             return self

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

@@ -1 +0,0 @@
-CHILD_DATA_VARIABLE_UUID = 'output_0'

mage_ai/data_preparation/models/block/dynamic/counter.py

@@ -1,14 +1,12 @@
 from typing import Any, Iterable, List, Optional, Sequence
 
 from mage_ai.data.tabular.reader import read_metadata
-from mage_ai.data_preparation.models.block.dynamic.constants import (
-    CHILD_DATA_VARIABLE_UUID,
-)
 from mage_ai.data_preparation.models.block.dynamic.utils import (
     is_dynamic_block,
     is_dynamic_block_child,
     should_reduce_output,
 )
+from mage_ai.data_preparation.models.constants import CHILD_DATA_VARIABLE_UUID
 from mage_ai.data_preparation.models.variables.cache import VariableAggregateCache
 from mage_ai.data_preparation.models.variables.constants import (
     VariableAggregateDataType,

mage_ai/data_preparation/models/block/outputs.py

@@ -343,7 +343,7 @@ def format_output_data(
                 columns=columns_to_display,
                 rows=[
                     list(row.values())
-                    for row in json.loads(data[columns_to_display].write_json(row_oriented=True))
+                    for row in json.loads(json.dumps(data[columns_to_display].to_dicts()))
                 ],
             ),
             resource_usage=resource_usage,

mage_ai/data_preparation/models/block/r/__init__.py

@@ -8,7 +8,10 @@ import pandas as pd
 import simplejson
 
 from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.data_preparation.models.constants import (
+    CHILD_DATA_VARIABLE_UUID,
+    BlockType,
+)
 from mage_ai.data_preparation.models.variables.constants import (
     DATAFRAME_CSV_FILE,
     VariableType,
@@ -81,7 +84,10 @@ def execute_r_code(
 
     if len(output_variable_objects) > 0:
         df = pd.read_csv(
-            os.path.join(output_variable_objects[0].variable_path, DATAFRAME_CSV_FILE)
+            os.path.join(
+                output_variable_objects[0].variable_path,
+                DATAFRAME_CSV_FILE
+            )
         )
     else:
         df = None
@@ -129,14 +135,19 @@ def __render_r_script(
         raise Exception(
             f"Block execution for {block.type} with R language is not supported.",
         )
-    template = template_env.get_template(BLOCK_TYPE_TO_EXECUTION_TEMPLATE[block.type])
+    template = template_env.get_template(
+        BLOCK_TYPE_TO_EXECUTION_TEMPLATE[block.type]
+    )
 
     output_variable_object = block.variable_object(
-        "output_0",
+        CHILD_DATA_VARIABLE_UUID,
         execution_partition=execution_partition,
     )
     os.makedirs(output_variable_object.variable_path, exist_ok=True)
-    output_path = os.path.join(output_variable_object.variable_path, DATAFRAME_CSV_FILE)
+    output_path = os.path.join(
+        output_variable_object.variable_path,
+        DATAFRAME_CSV_FILE
+    )
 
     global_vars_str = __render_global_vars(global_vars=global_vars)
 

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

@@ -150,6 +150,7 @@ def execute_sql_code(
     interpolate_vars_options = dict(
         block=block,
         dynamic_block_index=dynamic_block_index,
+        execution_partition=execution_partition,
         global_vars=global_vars,
     )
 
@@ -368,6 +369,7 @@ def execute_sql_code(
                     loader,
                     block,
                     query_string,
+                    disable_query_preprocessing=disable_query_preprocessing,
                     configuration=configuration,
                     should_query=should_query,
                 )

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

@@ -67,9 +67,17 @@ def execute_raw_sql(
     query_string: str,
     configuration: Dict = None,
     should_query: bool = False,
+    disable_query_preprocessing: bool = False,
 ) -> List[Any]:
     if configuration is None:
         configuration = {}
+
+    if disable_query_preprocessing:
+        return loader.execute_query_raw(
+            query_string,
+            configuration=configuration,
+        )
+
     queries = []
     fetch_query_at_indexes = []
 

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

@@ -232,7 +232,8 @@ def interpolate_vars(
     content: str,
     global_vars: Dict = None,
     block=None,
-    dynamic_block_index: int = None
+    dynamic_block_index: int = None,
+    execution_partition: str = None,
 ) -> str :
     if not content:
         return content
@@ -241,7 +242,10 @@ def interpolate_vars(
 
     if block:
         content = block.interpolate_content(
-            content, variables=global_vars, dynamic_block_index=dynamic_block_index,
+            content,
+            variables=global_vars,
+            dynamic_block_index=dynamic_block_index,
+            execution_partition=execution_partition,
         )
 
     return Template(

mage_ai/data_preparation/models/constants.py

@@ -27,6 +27,7 @@ PIPELINES_FOLDER = 'pipelines'
 PREFERENCES_FILE = '.preferences.yaml'
 REPO_CONFIG_FILE = 'metadata.yaml'
 VARIABLE_DIR = '.variables'
+CHILD_DATA_VARIABLE_UUID = 'output_0'
 
 PIPELINE_RUN_STATUS_LAST_RUN_FAILED = 'last_run_failed'
 
@@ -81,6 +82,7 @@ class BlockColor(StrEnum):
 
 
 class CallbackStatus(StrEnum):
+    CANCELLED = 'cancelled'
     FAILURE = 'failure'
     SUCCESS = 'success'
 
@@ -125,6 +127,7 @@ BLOCK_LANGUAGE_TO_FILE_EXTENSION = {
 
 
 CALLBACK_STATUSES = [
+    CallbackStatus.CANCELLED,
     CallbackStatus.FAILURE,
     CallbackStatus.SUCCESS,
 ]

mage_ai/data_preparation/models/pipeline.py

@@ -2331,7 +2331,7 @@ class Pipeline:
     ):
         # Introduce a small delay to prevent multiple changes from generating
         # identical timestamps for the pipeline YAML file
-        time.sleep(0.0004)
+        time.sleep(0.0005)
 
         blocks_current = sorted([b.uuid for b in self.blocks_by_uuid.values()])
 

mage_ai/data_preparation/storage/local_storage.py

@@ -121,7 +121,10 @@ class LocalStorage(BaseStorage):
         return pd.read_parquet(file_path, engine='pyarrow')
 
     def read_polars_parquet(self, file_path: str, **kwargs) -> pl.DataFrame:
-        return pl.read_parquet(file_path, use_pyarrow=True)
+        try:
+            return pl.read_parquet(file_path, use_pyarrow=True)
+        except Exception:
+            return pl.read_parquet(file_path)
 
     def write_csv(self, df: pd.DataFrame, file_path: str) -> None:
         File.create_parent_directories(file_path)