mage-ai 0.9.76__py3-none-any.whl → 0.9.78__py3-none-any.whl

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

@@ -215,7 +215,7 @@ def check_connection(git, remote_url: str) -> None:
 
 
 async def check_connection_async(git, remote_name: str) -> None:
-    proc = git.ls_remote(remote_name, as_process=True)
+    proc = git.ls_remote(remote_name, 'HEAD', as_process=True)
 
     await poll_process_with_timeout(
         proc,
@@ -228,7 +228,7 @@ async def check_connection_async(git, remote_name: str) -> None:
 
 
 async def validate_authentication_for_remote_url(git, remote_url: str) -> None:
-    proc = git.ls_remote(remote_url, as_process=True)
+    proc = git.ls_remote(remote_url, 'HEAD', as_process=True)
 
     await poll_process_with_timeout(
         proc,

mage_ai/data_preparation/logging/logger_manager.py

@@ -139,15 +139,31 @@ class LoggerManager:
                          str_to_timedelta(log_retention_period)).strftime(
                             format='%Y%m%dT%H%M%S')
 
+        from mage_ai.data_preparation.models.constants import PipelineType
+        from mage_ai.data_preparation.models.pipeline import Pipeline
         if self.pipeline_uuid is None:
-            from mage_ai.data_preparation.models.pipeline import Pipeline
-
             pipeline_uuids = Pipeline.get_all_pipelines(self.repo_path)
         else:
             pipeline_uuids = [self.pipeline_uuid]
 
         for pipeline_uuid in pipeline_uuids:
             print(f'Removing old logs from pipeline {pipeline_uuid}')
+            # Skip cleaning up logs for streaming pipelines
+            skip_clean_up_logs = False
+            try:
+                p = Pipeline.get(pipeline_uuid)
+                if p.type == PipelineType.STREAMING:
+                    print(f'Pipeline {pipeline_uuid} is a Streaming pipeline. '
+                          'Skip cleaning up logs.')
+                    skip_clean_up_logs = True
+            except Exception as e:
+                print(f'Fail to load pipeline {pipeline_uuid} due to error {e}. '
+                      'Skip cleaning up logs.')
+                skip_clean_up_logs = True
+
+            if skip_clean_up_logs:
+                continue
+
             pipeline_log_path = self.get_log_filepath_prefix(pipeline_uuid=pipeline_uuid)
             dirs = self.storage.listdir(pipeline_log_path)
             for dirname in dirs:

mage_ai/data_preparation/models/block/__init__.py

@@ -4464,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/outputs.py

@@ -343,7 +343,13 @@ 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(
+                        simplejson.dumps(
+                            data[columns_to_display].to_dicts(),
+                            default=encode_complex,
+                            ignore_nan=True,
+                        )
+                    )
                 ],
             ),
             resource_usage=resource_usage,

mage_ai/data_preparation/models/constants.py

@@ -82,6 +82,7 @@ class BlockColor(StrEnum):
 
 
 class CallbackStatus(StrEnum):
+    CANCELLED = 'cancelled'
     FAILURE = 'failure'
     SUCCESS = 'success'
 
@@ -126,6 +127,7 @@ BLOCK_LANGUAGE_TO_FILE_EXTENSION = {
 
 
 CALLBACK_STATUSES = [
+    CallbackStatus.CANCELLED,
     CallbackStatus.FAILURE,
     CallbackStatus.SUCCESS,
 ]

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)

mage_ai/io/config.py

@@ -66,6 +66,7 @@ class ConfigKey(StrEnum):
     MONGODB_PORT = 'MONGODB_PORT'
     MONGODB_USER = 'MONGODB_USER'
 
+    MSSQL_AUTHENTICATION = 'MSSQL_AUTHENTICATION'
     MSSQL_DATABASE = 'MSSQL_DATABASE'
     MSSQL_DRIVER = 'MSSQL_DRIVER'
     MSSQL_HOST = 'MSSQL_HOST'

mage_ai/io/mssql.py

@@ -38,12 +38,14 @@ class MSSQL(BaseSQL):
         host: str,
         password: str,
         user: str,
+        authentication: str = None,
         schema: str = None,
         port: int = 1433,
         verbose: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(
+            authentication=authentication,
             database=database,
             server=host,
             user=user,
@@ -61,15 +63,19 @@ class MSSQL(BaseSQL):
         database = self.settings['database']
         username = self.settings['user']
         password = self.settings['password']
-        return (
-            f'DRIVER={{{driver}}};'
-            f'SERVER={server};'
-            f'DATABASE={database};'
-            f'UID={username};'
-            f'PWD={password};'
-            'ENCRYPT=yes;'
-            'TrustServerCertificate=yes;'
-        )
+        conn_opts = [
+            f'DRIVER={{{driver}}}',
+            f'SERVER={server}',
+            f'DATABASE={database}',
+            f'UID={username}',
+            f'PWD={password}',
+            'ENCRYPT=yes',
+            'TrustServerCertificate=yes',
+        ]
+        if self.settings.get('authentication'):
+            authentication = self.settings.get('authentication')
+            conn_opts.append(f'Authentication={authentication}')
+        return ';'.join(conn_opts)
 
     def default_schema(self) -> str:
         return self.settings.get('schema') or 'dbo'
@@ -318,6 +324,7 @@ class MSSQL(BaseSQL):
     @classmethod
     def with_config(cls, config: BaseConfigLoader) -> 'MSSQL':
         return cls(
+            authentication=config[ConfigKey.MSSQL_AUTHENTICATION],
             database=config[ConfigKey.MSSQL_DATABASE],
             schema=config[ConfigKey.MSSQL_SCHEMA],
             driver=config[ConfigKey.MSSQL_DRIVER],

mage_ai/io/postgres.py

@@ -329,18 +329,21 @@ class Postgres(BaseSQL):
                 return simplejson.dumps(
                     val,
                     default=encode_complex,
+                    ensure_ascii=False,
                     ignore_nan=True,
                 )
             elif type(val) is list and len(val) >= 1 and type(val[0]) is dict:
                 return simplejson.dumps(
                     val,
                     default=encode_complex,
+                    ensure_ascii=False,
                     ignore_nan=True,
                 )
             elif not use_insert_command and type(val) is list:
                 return self._clean_array_value(simplejson.dumps(
                     val,
                     default=encode_complex,
+                    ensure_ascii=False,
                     ignore_nan=True,
                 ))
             return val

mage_ai/orchestration/db/migrations/versions/39d36f1dab73_create_genericjob.py

@@ -0,0 +1,47 @@
+"""Create GenericJob
+
+Revision ID: 39d36f1dab73
+Revises: 0227396a216c
+Create Date: 2025-08-27 20:45:29.756869
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '39d36f1dab73'
+down_revision = '0227396a216c'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('generic_job',
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
+    sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
+    sa.Column('job_type', sa.Enum('CANCEL_PIPELINE_RUN', name='jobtype'), nullable=False),
+    sa.Column('status', sa.Enum('INITIAL', 'QUEUED', 'RUNNING', 'COMPLETED', 'FAILED', name='jobstatus'), nullable=True),
+    sa.Column('payload', sa.JSON(), nullable=True),
+    sa.Column('extra_metadata', sa.JSON(), nullable=True),
+    sa.Column('started_at', sa.DateTime(timezone=True), nullable=True),
+    sa.Column('completed_at', sa.DateTime(timezone=True), nullable=True),
+    sa.Column('repo_path', sa.String(length=255), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    with op.batch_alter_table('generic_job', schema=None) as batch_op:
+        batch_op.create_index(batch_op.f('ix_generic_job_job_type'), ['job_type'], unique=False)
+        batch_op.create_index(batch_op.f('ix_generic_job_status'), ['status'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('generic_job', schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f('ix_generic_job_status'))
+        batch_op.drop_index(batch_op.f('ix_generic_job_job_type'))
+
+    op.drop_table('generic_job')
+    # ### end Alembic commands ###

mage_ai/orchestration/db/models/oauth.py

@@ -882,4 +882,5 @@ class Oauth2AccessToken(BaseModel):
     def is_valid(self) -> bool:
         return self.token and \
             self.expires and \
-            self.expires >= datetime.utcnow().replace(tzinfo=self.expires.tzinfo)
+            self.expires >= datetime.utcnow().replace(tzinfo=self.expires.tzinfo) and \
+            self.user is not None

mage_ai/orchestration/db/models/schedules.py

@@ -1804,6 +1804,111 @@ class BlockRun(BlockRunProjectPlatformMixin, BaseModel):
         )
 
 
+class GenericJob(BaseModel):
+    class JobStatus(StrEnum):
+        INITIAL = 'initial'
+        QUEUED = 'queued'
+        RUNNING = 'running'
+        COMPLETED = 'completed'
+        FAILED = 'failed'
+
+    class JobType(StrEnum):
+        CANCEL_PIPELINE_RUN = 'cancel_pipeline_run'
+
+    job_type = Column(Enum(JobType), nullable=False, index=True)
+    status = Column(Enum(JobStatus), default=JobStatus.INITIAL, index=True)
+    payload = Column(JSON)  # Job-specific data (e.g., pipeline_run_id)
+    extra_metadata = Column(JSON)  # Additional metadata for the job
+    started_at = Column(DateTime(timezone=True))
+    completed_at = Column(DateTime(timezone=True))
+    repo_path = Column(String(255))
+
+    @classproperty
+    def repo_query(cls):
+        return cls.query.filter(
+            or_(
+                GenericJob.repo_path == get_repo_path(),
+                GenericJob.repo_path.is_(None),
+            )
+        )
+
+    @classmethod
+    @safe_db_query
+    def enqueue_cancel_pipeline_run(
+        cls,
+        pipeline_run_id: int,
+        cancelled_block_run_ids: List,
+        repo_path: str = None,
+    ) -> 'GenericJob':
+        """Enqueue a job to cancel a pipeline run."""
+        if repo_path is None:
+            repo_path = get_repo_path()
+
+        job = cls.create(
+            job_type=cls.JobType.CANCEL_PIPELINE_RUN,
+            payload=dict(
+                pipeline_run_id=pipeline_run_id,
+                cancelled_block_run_ids=cancelled_block_run_ids,
+            ),
+            repo_path=repo_path,
+            status=cls.JobStatus.INITIAL
+        )
+        return job
+
+    @classmethod
+    @safe_db_query
+    def get_jobs_with_initial_status(
+        cls,
+        limit: int = 10,
+        repo_path: str = None,
+    ) -> List['GenericJob']:
+        """Get up to N jobs with INITIAL status for processing."""
+        query = cls.repo_query.filter(
+            cls.status == cls.JobStatus.INITIAL,
+        )
+
+        if repo_path:
+            query = query.filter(cls.repo_path == repo_path)
+
+        jobs = query.order_by(cls.created_at.asc()).limit(limit).all()
+
+        return jobs
+
+    @safe_db_query
+    def mark_completed(self):
+        """Mark the job as completed successfully."""
+        self.update(
+            status=self.JobStatus.COMPLETED,
+            completed_at=datetime.now(tz=pytz.UTC)
+        )
+
+    @safe_db_query
+    def mark_failed(self, metadata: dict = None):
+        """Mark the job as failed."""
+        self.update(
+            status=self.JobStatus.FAILED,
+            completed_at=datetime.now(tz=pytz.UTC),
+            extra_metadata=metadata or {}
+        )
+
+    @safe_db_query
+    def mark_queued(self):
+        """Mark the job as queued."""
+        self.update(
+            status=self.JobStatus.QUEUED,
+        )
+
+    @safe_db_query
+    def mark_running(self):
+        """Mark the job as queued."""
+        self.update(
+            status=self.JobStatus.RUNNING,
+        )
+
+    def __repr__(self):
+        return f'GenericJob(id={self.id}, type={self.job_type}, status={self.status})'
+
+
 class EventMatcher(BaseModel):
     class EventType(StrEnum):
         AWS_EVENT = 'aws_event'

mage_ai/orchestration/job_manager.py

@@ -8,6 +8,7 @@ class JobType(StrEnum):
     BLOCK_RUN = 'block_run'
     PIPELINE_RUN = 'pipeline_run'
     INTEGRATION_STREAM = 'integration_stream'
+    GENERIC_JOB = 'generic_job'
 
 
 class JobManager:
@@ -62,6 +63,19 @@ class JobManager:
             logging_tags=logging_tags,
         )
 
+    def has_generic_job(
+        self,
+        generic_job_id: int,
+        logger=None,
+        logging_tags: Dict = None,
+    ) -> bool:
+        job_id = self.__job_id(JobType.GENERIC_JOB, generic_job_id)
+        return self.queue.has_job(
+            job_id,
+            logger=logger,
+            logging_tags=logging_tags,
+        )
+
     def kill_block_run_job(self, block_run_id):
         print(f'Kill block run id: {block_run_id}')
         job_id = self.__job_id(JobType.BLOCK_RUN, block_run_id)
@@ -78,6 +92,11 @@ class JobManager:
         job_id = self.__job_id(JobType.INTEGRATION_STREAM, id)
         return self.queue.kill_job(job_id)
 
+    def kill_generic_job(self, generic_job_id):
+        print(f'Kill generic job id: {generic_job_id}')
+        job_id = self.__job_id(JobType.GENERIC_JOB, generic_job_id)
+        return self.queue.kill_job(job_id)
+
     def start(self):
         self.queue.start()
 

mage_ai/orchestration/notification/sender.py

@@ -69,7 +69,7 @@ class NotificationSender:
 
         if self.config.teams_config is not None and self.config.teams_config.is_valid:
             try:
-                send_teams_message(self.config.teams_config, summary)
+                send_teams_message(self.config.teams_config, summary, title)
             except Exception:
                 traceback.print_exc()
 
@@ -87,7 +87,7 @@ class NotificationSender:
 
         if self.config.google_chat_config is not None and self.config.google_chat_config.is_valid:
             try:
-                send_google_chat_message(self.config.google_chat_config, summary)
+                send_google_chat_message(self.config.google_chat_config, summary, title)
             except Exception:
                 traceback.print_exc()