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

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -32,6 +32,7 @@ from mage_ai.orchestration.db.models.schedules import (
     Backfill,
     BlockRun,
     EventMatcher,
+    GenericJob,
     PipelineRun,
     PipelineSchedule,
 )
@@ -928,6 +929,81 @@ class PipelineScheduler:
             self.memory_usage_failure(tags=tags)
 
 
+def on_pipeline_run_cancelled(
+    job_id: int,
+    pipeline_run_id: int,
+    cancelled_block_run_ids: List,
+):
+    job = GenericJob.query.get(job_id)
+    if job.status not in [
+        GenericJob.JobStatus.INITIAL,
+        GenericJob.JobStatus.QUEUED,
+        GenericJob.JobStatus.RUNNING,
+    ]:
+        return
+    job.mark_running()
+
+    if not pipeline_run_id:
+        job.mark_failed(
+            metadata=dict(
+                error_message='Pipeline run id is not provided in the job.',
+            ),
+        )
+        return
+    if not cancelled_block_run_ids:
+        job.mark_failed(
+            metadata=dict(
+                error_message='Cancelled block run ids are not provided in the job.',
+            ),
+        )
+        return
+
+    try:
+        # Run callback blocks for cancelled block runs
+        pipeline_run = PipelineRun.query.get(pipeline_run_id)
+        if not pipeline_run:
+            raise Exception(f'Fail to retrieve pipeline run with id {pipeline_run_id}')
+        pipeline_schedule = pipeline_run.pipeline_schedule
+        pipeline = get_pipeline_from_platform(
+            pipeline_run.pipeline_uuid,
+            repo_path=pipeline_schedule.repo_path if pipeline_schedule else None,
+        )
+        if pipeline_run.status != PipelineRun.PipelineRunStatus.CANCELLED:
+            job.mark_failed(
+                metadata=dict(
+                    error_message=f'Pipeline run {pipeline_run_id} is not in cancelled status.',
+                ),
+            )
+            return
+        for block_run_id in cancelled_block_run_ids:
+            block_run = BlockRun.query.get(block_run_id)
+            if not block_run:
+                continue
+            if block_run.status != BlockRun.BlockRunStatus.CANCELLED:
+                continue
+            ExecutorFactory.get_block_executor(
+                pipeline,
+                block_run.block_uuid,
+                block_run_id=block_run.id,
+                execution_partition=pipeline_run.execution_partition,
+                executor_type=ExecutorType.LOCAL_PYTHON,
+            ).execute_callback(
+                callback='on_cancelled',
+                global_vars=pipeline_run.get_variables(),
+                logging_tags=None,
+                pipeline_run=pipeline_run,
+                block_run_id=block_run_id,
+            )
+    except Exception as e:
+        job.mark_failed(
+            metadata=dict(
+                error_message=f'Failed to run pipeline run cancellation hook: {e}.',
+            ),
+        )
+        return
+    job.mark_completed()
+
+
 def run_integration_streams(
     streams: List[Dict],
     *args,
@@ -1406,8 +1482,9 @@ def cancel_block_runs_and_jobs(
             block_runs_to_cancel.append(b)
         if b.status == BlockRun.BlockRunStatus.RUNNING:
             running_blocks.append(b)
+    cancelled_block_run_ids = [b.id for b in block_runs_to_cancel]
     BlockRun.batch_update_status(
-        [b.id for b in block_runs_to_cancel],
+        cancelled_block_run_ids,
         BlockRun.BlockRunStatus.CANCELLED,
     )
 
@@ -1435,6 +1512,18 @@ def cancel_block_runs_and_jobs(
         for b in running_blocks:
             job_manager.kill_block_run_job(b.id)
 
+    GenericJob.enqueue_cancel_pipeline_run(pipeline_run.id, cancelled_block_run_ids)
+
+
+def kill_cancelled_runs(pipeline_run_id: int, cancelled_block_run_ids: List):
+    job_manager = get_job_manager()
+    if pipeline_run_id and job_manager.has_pipeline_run_job(pipeline_run_id):
+        job_manager.kill_pipeline_run_job(pipeline_run_id)
+    if cancelled_block_run_ids:
+        for block_run_id in cancelled_block_run_ids:
+            if job_manager.has_block_run_job(block_run_id):
+                job_manager.kill_block_run_job(block_run_id)
+
 
 def check_sla():
     repo_pipelines = set(
@@ -1732,6 +1821,13 @@ def schedule_all():
             logger.exception(f'Failed to schedule {r}')
             traceback.print_exc()
             continue
+
+    try:
+        schedule_generic_jobs()
+    except Exception as e:
+        logger.exception(f'Failed to schedule generic jobs {e}')
+        traceback.print_exc()
+
     job_manager = get_job_manager()
     if job_manager is None:
         logger.info('Job manager is None.')
@@ -1912,3 +2008,52 @@ def sync_schedules(pipeline_uuids: List[str]):
             trigger_configs.append(pipeline_trigger)
 
     PipelineSchedule.create_or_update_batch(trigger_configs)
+
+
+def schedule_generic_jobs():
+    """
+    Schedule generic jobs that are in INITIAL status.
+
+    This method schedules jobs like cancel_pipeline_run that have been enqueued
+    by the server process and are waiting to be executed by the job manager.
+    """
+    try:
+        # Get all jobs with INITIAL status across all repositories
+        jobs = GenericJob.get_jobs_with_initial_status()
+        job_manager = get_job_manager()
+
+        if not jobs:
+            return
+
+        logger.info(f'Scheduling {len(jobs)} generic jobs')
+
+        # Schedule each job with the job manager
+        for job in jobs:
+            try:
+                if job.job_type == GenericJob.JobType.CANCEL_PIPELINE_RUN:
+                    pipeline_run_id = job.payload.get('pipeline_run_id')
+                    cancelled_block_run_ids = job.payload.get('cancelled_block_run_ids')
+                    try:
+                        kill_cancelled_runs(pipeline_run_id, cancelled_block_run_ids)
+                    except Exception:
+                        pass
+                    # Add job to job manager for execution
+                    job_manager.add_job(
+                        JobType.GENERIC_JOB,
+                        job.id,
+                        on_pipeline_run_cancelled,
+                        # args
+                        job.id,
+                        pipeline_run_id,
+                        cancelled_block_run_ids,
+                    )
+                    job.mark_queued()
+                    logger.info(f'Scheduled generic job {job.id} ({job.job_type})')
+                else:
+                    raise ValueError(f"Unknown job type: {job.job_type}")
+            except Exception as e:
+                error_message = f"Failed to schedule generic job: {str(e)}"
+                logger.exception(f'Failed to schedule generic job {job.id}: {error_message}')
+                job.mark_failed({'error': error_message})
+    except Exception as e:
+        logger.exception(f'Failed to schedule generic jobs: {e}')

mage_ai/orchestration/queue/config.py

@@ -1,8 +1,11 @@
+import traceback
 from dataclasses import dataclass
 
 from mage_ai.shared.config import BaseConfig
 from mage_ai.shared.enum import StrEnum
 
+DEFAULT_CONCURRENCY = 20
+
 
 class QueueType(StrEnum):
     CELERY = 'celery'
@@ -17,5 +20,12 @@ class ProcessQueueConfig(BaseConfig):
 @dataclass
 class QueueConfig(BaseConfig):
     queue_type: QueueType = QueueType.PROCESS
-    concurrency: int = 20
+    concurrency: int = DEFAULT_CONCURRENCY
     process_queue_config: ProcessQueueConfig = None
+
+    def __post_init__(self):
+        try:
+            self.concurrency = int(self.concurrency)
+        except Exception:
+            traceback.print_exc()
+            self.concurrency = DEFAULT_CONCURRENCY

mage_ai/orchestration/queue/process_queue.py

@@ -19,6 +19,7 @@ from mage_ai.settings import (
     HOSTNAME,
     REDIS_URL,
     SENTRY_DSN,
+    SENTRY_SERVER_NAME,
     SENTRY_TRACES_SAMPLE_RATE,
     SERVER_LOGGING_FORMAT,
     SERVER_VERBOSITY,
@@ -287,6 +288,7 @@ class Worker(mp.Process):
             sentry_sdk.init(
                 self.dsn,
                 traces_sample_rate=SENTRY_TRACES_SAMPLE_RATE,
+                server_name=SENTRY_SERVER_NAME,
             )
         initialize_new_relic()
 

mage_ai/server/api/base.py

@@ -1,12 +1,15 @@
 import asyncio
 import json
+import re
 import traceback
+import urllib.parse
 
 import dateutil.parser
 import simplejson
 import tornado.web
 
 from mage_ai.api.middleware import OAuthMiddleware
+from mage_ai.server.api.constants import PATH_TRAVERSAL_PATTERN
 from mage_ai.shared.parsers import encode_complex
 from mage_ai.shared.strings import camel_to_snake_case
 from mage_ai.usage_statistics.constants import EventNameType
@@ -127,11 +130,49 @@ class BaseApiHandler(BaseHandler, OAuthMiddleware):
         super().initialize(**kwargs)
         self.is_health_check = kwargs.get('is_health_check', False)
 
+    def is_safe_path(self, user_input):
+        """
+        Check if the user input is safe and doesn't contain path traversal sequences.
+        """
+        return re.match(PATH_TRAVERSAL_PATTERN, user_input) is not None
+
     def prepare(self):
         from mage_ai.server.server import latest_user_activity
 
         if not self.is_health_check:
             latest_user_activity.update_latest_activity()
+
+        # Validate the request path by decoding from bytes to string if necessary
+        decoded_path = self.request.path
+        if isinstance(decoded_path, bytes):
+            decoded_path = decoded_path.decode('utf-8')  # Decode only if it's a bytes object
+
+        decoded_path = urllib.parse.unquote(decoded_path)  # Decode URL-encoded characters
+        if not self.is_safe_path(decoded_path):
+            self.set_status(400)  # Bad Request
+            self.write("Error: Invalid path (path traversal detected)")
+            self.finish()
+            return
+
+        # Validate query parameters (if any)
+        for key, value in self.request.arguments.items():
+            # Decode each key and value from bytes to string if necessary
+            decoded_key = key
+            if isinstance(decoded_key, bytes):
+                decoded_key = decoded_key.decode('utf-8')
+
+            decoded_value = value[0]
+            if isinstance(decoded_value, bytes):
+                decoded_value = decoded_value.decode('utf-8')
+
+            decoded_value = urllib.parse.unquote(decoded_value)  # Decode URL-encoded characters
+            if not self.is_safe_path(decoded_value):
+                self.set_status(400)
+                self.write(
+                    f"Error: Invalid parameter value for '{decoded_key}' (path traversal detected)"
+                )
+                self.finish()
+                return
         super().prepare()
 
 

mage_ai/server/api/constants.py

@@ -3,4 +3,5 @@ ENDPOINTS_BYPASS_OAUTH_CHECK = []
 HEADER_API_KEY = 'X-API-KEY'
 HEADER_OAUTH_TOKEN = 'OAUTH-TOKEN'
 OAUTH_CLIENT_ID_KEY = 'OAUTH_CLIENT_ID'
+PATH_TRAVERSAL_PATTERN = r"^(?!.*(\.\.\/|\.\.\\)).*$"
 URL_PARAMETER_API_KEY = 'api_key'

mage_ai/server/constants.py

@@ -12,4 +12,4 @@ DATAFRAME_OUTPUT_SAMPLE_COUNT = 10
 # Dockerfile depends on it because it runs ./scripts/install_mage.sh and uses
 # the last line to determine the version to install.
 VERSION = \
-'0.9.76'
+'0.9.77'

mage_ai/server/scheduler_manager.py

@@ -12,6 +12,7 @@ from mage_ai.server.logger import Logger
 from mage_ai.services.newrelic import initialize_new_relic
 from mage_ai.settings import (
     SENTRY_DSN,
+    SENTRY_SERVER_NAME,
     SENTRY_TRACES_SAMPLE_RATE,
     SERVER_LOGGING_FORMAT,
     SERVER_VERBOSITY,
@@ -35,6 +36,7 @@ def run_scheduler():
         sentry_sdk.init(
             sentry_dsn,
             traces_sample_rate=SENTRY_TRACES_SAMPLE_RATE,
+            server_name=SENTRY_SERVER_NAME,
         )
     (enable_new_relic, application) = initialize_new_relic()
     try:

mage_ai/server/terminal_server.py

@@ -126,6 +126,9 @@ class TerminalWebsocketServer(terminado.TermSocket):
                             Entity.PROJECT,
                             get_project_uuid(),
                         )
+                    else:
+                        # If the user is deleted, set valid to False
+                        valid = False
             if not valid or is_disable_pipeline_edit_access():
                 return self.send_json_message(
                     ['stdout', f'{command[1]}\nUnauthorized access to the terminal.'])

mage_ai/settings/server.py

@@ -154,6 +154,7 @@ USE_UNIQUE_TERMINAL = os.getenv('USE_UNIQUE_TERMINAL', None)
 # Sentry Configuration
 SENTRY_DSN = os.getenv('SENTRY_DSN')
 SENTRY_TRACES_SAMPLE_RATE = os.getenv('SENTRY_TRACES_SAMPLE_RATE', 1.0)
+SENTRY_SERVER_NAME = os.getenv('SENTRY_SERVER_NAME', None)
 
 # New Relic Configuration
 ENABLE_NEW_RELIC = os.getenv('ENABLE_NEW_RELIC', False)
@@ -265,6 +266,7 @@ MAGE_SETTINGS_ENVIRONMENT_VARIABLES = [
     'USE_UNIQUE_TERMINAL',
     'SENTRY_DSN',
     'SENTRY_TRACES_SAMPLE_RATE',
+    'SENTRY_SERVER_NAME',
     'MAGE_PUBLIC_HOST',
     'SCHEDULER_TRIGGER_INTERVAL',
     'REQUIRE_USER_PERMISSIONS',

mage_ai/streaming/sources/kafka.py

@@ -91,7 +91,7 @@ class KafkaSource(BaseSource):
             api_version=self.config.api_version,
             auto_offset_reset=self.config.auto_offset_reset,
             max_partition_fetch_bytes=self.config.max_partition_fetch_bytes,
-            enable_auto_commit=True,
+            enable_auto_commit=False,
         )
         if self.config.security_protocol == SecurityProtocol.SSL:
             consumer_kwargs['security_protocol'] = SecurityProtocol.SSL
@@ -270,6 +270,7 @@ class KafkaSource(BaseSource):
                     message_values.append(message)
             if len(message_values) > 0:
                 handler(message_values)
+            self.consumer.commit()
 
     def test_connection(self):
         self.consumer._client.check_version(timeout=5)

mage_ai/tests/data_preparation/executors/test_block_executor.py

@@ -98,7 +98,7 @@ class BlockExecutorTest(BaseApiTestCase):
         self.block_executor._execute_conditional = MagicMock(return_value=True)
         self.block_executor._execute = MagicMock(return_value={'result': 'success'})
         # self.block.run_tests = MagicMock()
-        self.block_executor._execute_callback = MagicMock()
+        self.block_executor.execute_callback = MagicMock()
 
         result = self.block_executor.execute(
             analyze_outputs=analyze_outputs,
@@ -167,7 +167,7 @@ class BlockExecutorTest(BaseApiTestCase):
         #     update_tests=False,
         #     dynamic_block_uuid=dynamic_block_uuid,
         # )
-        self.block_executor._execute_callback.assert_called_with(
+        self.block_executor.execute_callback.assert_called_with(
             'on_success',
             block_run_id=None,
             callback_kwargs=dict(retry=dict(attempts=1)),
@@ -245,7 +245,7 @@ class BlockExecutorTest(BaseApiTestCase):
         self.block.callback_blocks = [MagicMock(), MagicMock()]
         self.block.callback_block = MagicMock()
 
-        self.block_executor._execute_callback(
+        self.block_executor.execute_callback(
             callback='on_success',
             global_vars=dict(retry=dict(attempts=1)),
             logging_tags={},

mage_ai/tests/data_preparation/models/test_variable.py

@@ -180,6 +180,7 @@ class VariableTest(DBTestCase):
                 [1, 'test'],
                 [2, 'test2'],
             ],
+            orient="row",
             schema=['col1', 'col2'],
         )
         df2 = pl.DataFrame(
@@ -187,6 +188,7 @@ class VariableTest(DBTestCase):
                 [1, 'test', 3.123, 41414123123124],
                 [2, 'test2', 4.321, 12111111],
             ],
+            orient="row",
             schema=['col1', 'col2', 'col3', 'col4'],
         )
         df2 = df2.cast({'col4': pl.Int64})

mage_ai/tests/io/create_table/test_postgresql.py

@@ -33,14 +33,14 @@ class TestTablePostgres(DBTestCase):
         test_cases = [
             [],
             [123],
-            [['abc', 'def']],
+            [['àabc', 'deèéf']],
             [['08:00', '12:00'], ['15:00', '20:00']],
             [['08:00', '12:00'], []],
         ]
         expected = [
             '{}',
             '{123}',
-            '{{"abc", "def"}}',
+            '{{"àabc", "deèéf"}}',
             '{{"08:00", "12:00"}, {"15:00", "20:00"}}',
             '{{"08:00", "12:00"}, {}}',
         ]
@@ -49,6 +49,7 @@ class TestTablePostgres(DBTestCase):
                 simplejson.dumps(
                     val,
                     default=encode_complex,
+                    ensure_ascii=False,
                     ignore_nan=True,
                 )
             )

mage_ai/tests/orchestration/notification/test_sender.py

@@ -112,9 +112,11 @@ class NotificationSenderTests(DBTestCase):
             f'`{pipeline_run.pipeline_schedule.name}` '
             f'at execution time `{pipeline_run.execution_date}`. Error: None'
         )
+        title = 'Failed to run Mage pipeline test_pipeline'
         mock_send_teams_message.assert_called_once_with(
             notification_config.teams_config,
             message,
+            title
         )
 
     @patch('mage_ai.orchestration.notification.sender.send_teams_message')
@@ -131,9 +133,11 @@ class NotificationSenderTests(DBTestCase):
             f'`{pipeline_run.pipeline_schedule.name}` '
             f'at execution time `{pipeline_run.execution_date}`.'
         )
+        title = 'Successfully ran Pipeline test_pipeline'
         mock_send_teams_message.assert_called_once_with(
             notification_config.teams_config,
             message,
+            title
         )
 
     @patch('mage_ai.orchestration.notification.sender.send_teams_message')
@@ -162,4 +166,4 @@ class NotificationSenderTests(DBTestCase):
             notification_config.opsgenie_config,
             message=ANY,
             description=ANY,
-        )
+        )

mage_ai/tests/streaming/sources/test_kafka.py

@@ -33,7 +33,7 @@ class KafkaTests(TestCase):
             api_version='0.10.2',
             auto_offset_reset='latest',
             max_partition_fetch_bytes=1048576,
-            enable_auto_commit=True,
+            enable_auto_commit=False,
         )
 
     def test_init_client_with_topics(self):
@@ -56,7 +56,7 @@ class KafkaTests(TestCase):
             api_version='0.10.2',
             auto_offset_reset='latest',
             max_partition_fetch_bytes=1048576,
-            enable_auto_commit=True,
+            enable_auto_commit=False,
         )
 
     def test_init_client_with_missing_topic_or_topics(self):