mage-ai 0.9.69__py3-none-any.whl → 0.9.70__py3-none-any.whl

mage_ai/data_preparation/models/block/remote/models.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from typing import Any, List
+
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class RemoteBlock(BaseDataClass):
+    block_uuid: str
+    execution_partition: str = None
+    pipeline_uuid: str = None
+    repo_path: str = None
+
+    def __post_init__(self):
+        self._block = None
+        self._pipeline = None
+
+    @property
+    def pipeline(self):
+        from mage_ai.data_preparation.models.pipeline import Pipeline
+
+        if self._pipeline:
+            return self._pipeline
+
+        self._pipeline = Pipeline.get(
+            self.pipeline_uuid,
+            all_projects=True,
+            check_if_exists=False,
+            repo_path=self.repo_path,
+            use_repo_path=True if self.repo_path else False,
+        )
+
+        return self._pipeline
+
+    @property
+    def block(self):
+        if self._block:
+            return self._block
+
+        self._block = self.pipeline.get_block(self.block_uuid)
+
+        return self._block
+
+    @property
+    def variable_uuids(self):
+        return self.block.output_variables(execution_partition=self.execution_partition)
+
+    def get_outputs(self, **kwargs) -> List[Any]:
+        arr = []
+        for variable_uuid in self.variable_uuids:
+            output = self.pipeline.get_block_variable(
+                self.block_uuid,
+                variable_uuid,
+                partition=self.execution_partition,
+            )
+            arr.append(output)
+
+        return arr

mage_ai/data_preparation/models/block/utils.py

@@ -29,6 +29,7 @@ from mage_ai.data_preparation.models.block.dynamic.utils import (
 from mage_ai.data_preparation.models.block.dynamic.utils import (
     should_reduce_output as should_reduce_output_original,
 )
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.constants import (
     DATAFRAME_ANALYSIS_MAX_COLUMNS,
     BlockType,
@@ -393,6 +394,7 @@ def fetch_input_variables(
     dynamic_block_flags: List[DynamicBlockFlag] = None,
     metadata: Dict = None,
     upstream_block_uuids_override: List[str] = None,
+    current_block=None,
 ) -> Tuple[List, List, List]:
     """
     Fetches the input variables for a block.
@@ -510,6 +512,22 @@ def fetch_input_variables(
             if BlockType.GLOBAL_DATA_PRODUCT == upstream_block.type:
                 global_data_product = upstream_block.get_global_data_product()
                 input_vars[idx] = global_data_product.get_outputs()
+
+                mds = {}
+                variable_uuids = upstream_block.output_variables(
+                    execution_partition=execution_partition,
+                )
+                for variable_uuid in variable_uuids:
+                    md = pipeline.get_block_variable(
+                        upstream_block_uuid,
+                        variable_uuid,
+                        partition=execution_partition,
+                    )
+                    if isinstance(md, dict):
+                        mds.update(md)
+
+                kwargs_vars.append(mds)
+
                 continue
 
             # Block output variables for upstream_block_uuid
@@ -743,6 +761,26 @@ def fetch_input_variables(
         __uuid='output_variables'
     )
 
+    if kwargs_vars:
+        kwargs_vars2 = []
+
+        remote_blocks_output = []
+        for kwargs in kwargs_vars:
+            for remote_block_dict in kwargs.get('remote_blocks', []):
+                # Global data products only need the remote block information, not the output
+                if current_block and BlockType.GLOBAL_DATA_PRODUCT == current_block.type:
+                    output = remote_block_dict
+                else:
+                    output = RemoteBlock.load(**remote_block_dict).get_outputs()
+            remote_blocks_output.append(output)
+
+        for kwargs in kwargs_vars:
+            if kwargs.get('remote_blocks'):
+                kwargs['remote_blocks'] = remote_blocks_output
+            kwargs_vars2.append(kwargs)
+
+        kwargs_vars = kwargs_vars2
+
     return input_vars, kwargs_vars, upstream_block_uuids_final
 
 

mage_ai/data_preparation/models/constants.py

@@ -24,6 +24,8 @@ PREFERENCES_FILE = '.preferences.yaml'
 REPO_CONFIG_FILE = 'metadata.yaml'
 VARIABLE_DIR = '.variables'
 
+PIPELINE_RUN_STATUS_LAST_RUN_FAILED = 'last_run_failed'
+
 
 class AIMode(str, Enum):
     OPEN_AI = 'open_ai'

mage_ai/data_preparation/models/global_data_product/__init__.py

@@ -75,6 +75,18 @@ class GlobalDataProduct:
 
         return self._object
 
+    def get_blocks(self) -> List:
+        arr = []
+
+        if not self.settings or len(self.settings) == 0:
+            return arr
+
+        if GlobalDataProductObjectType.PIPELINE == self.object_type:
+            for block_uuid in self.settings.keys():
+                arr.append(self.pipeline.get_block(block_uuid))
+
+        return arr
+
     def get_outputs(self, from_notebook: bool = False, global_vars: Dict = None, **kwargs) -> Dict:
         data = {}
 

mage_ai/data_preparation/models/pipeline.py

@@ -7,7 +7,7 @@ import tempfile
 import zipfile
 from datetime import datetime, timezone
 from io import BytesIO
-from typing import Any, Callable, Dict, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import aiofiles
 import pytz
@@ -99,6 +99,7 @@ class Pipeline:
         self.description = description
         self.executor_config = dict()
         self.executor_type = None
+        self._rendered_executor_type = None  # Render template variables
         self.extensions = {}
         self.name = None
         self.notification_config = dict()
@@ -1768,10 +1769,13 @@ class Pipeline:
     def get_executable_blocks(self):
         return [b for b in self.blocks_by_uuid.values() if b.executable]
 
-    def get_executor_type(self) -> str:
-        if self.executor_type:
-            return Template(self.executor_type).render(**get_template_vars())
-        return self.executor_type
+    def get_executor_type(self) -> Optional[str]:
+        if self._rendered_executor_type is None:
+            if self.executor_type:
+                return Template(self.executor_type).render(**get_template_vars())
+            else:
+                self._rendered_executor_type = self.executor_type
+        return self._rendered_executor_type
 
     def has_block(self, block_uuid: str, block_type: str = None, extension_uuid: str = None):
         if extension_uuid:

mage_ai/data_preparation/models/triggers/__init__.py

@@ -206,6 +206,7 @@ def add_or_update_trigger_for_pipeline_and_persist(
     trigger: Trigger,
     pipeline_uuid: str,
     update_only_if_exists: bool = False,
+    old_trigger_name: str = None,
 ) -> Dict:
     trigger_configs_by_name = get_trigger_configs_by_name(pipeline_uuid)
 
@@ -214,13 +215,14 @@ def add_or_update_trigger_for_pipeline_and_persist(
     have, so we need to set "envs" on the updated trigger if it already exists.
     Otherwise, it will get overwritten when updating the trigger in code.
     """
-    existing_trigger = trigger_configs_by_name.get(trigger.name)
+    trigger_name = trigger.name if old_trigger_name is None else old_trigger_name
+    existing_trigger = trigger_configs_by_name.get(trigger_name)
     if existing_trigger is not None:
         trigger.envs = existing_trigger.get('envs', [])
     elif update_only_if_exists:
         return None
 
-    trigger_configs_by_name[trigger.name] = trigger.to_dict()
+    trigger_configs_by_name[trigger_name] = trigger.to_dict()
     yaml_config = dict(triggers=list(trigger_configs_by_name.values()))
     content = yaml.safe_dump(yaml_config)
     trigger_file_path = get_triggers_file_path(pipeline_uuid)

mage_ai/data_preparation/storage/local_storage.py

@@ -92,12 +92,18 @@ class LocalStorage(BaseStorage):
             os.makedirs(dirname, exist_ok=True)
 
         with open(file_path, 'w') as file:
-            simplejson.dump(
-                data,
-                file,
-                default=encode_complex,
-                ignore_nan=True,
-            )
+            try:
+                simplejson.dump(
+                    data,
+                    file,
+                    default=encode_complex,
+                    ignore_nan=True,
+                )
+            except ValueError as err:
+                if is_debug():
+                    raise err
+                else:
+                    print(f'[ERROR] LocalStorage.write_json_file: {err}')
 
     async def write_json_file_async(self, file_path: str, data) -> None:
         async with aiofiles.open(file_path, mode='w') as file:

mage_ai/orchestration/db/migrations/versions/42a14d6143f1_update_token_column_type.py

@@ -0,0 +1,54 @@
+"""Update token column type
+
+Revision ID: 42a14d6143f1
+Revises: b9a2d6d0a2c7
+Create Date: 2024-04-12 15:19:52.639580
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '42a14d6143f1'
+down_revision = 'b9a2d6d0a2c7'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('oauth2_access_token', schema=None) as batch_op:
+        batch_op.alter_column(
+            'token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+        batch_op.alter_column(
+            'refresh_token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('oauth2_access_token', schema=None) as batch_op:
+        batch_op.alter_column(
+            'token',
+            existing_type=sa.Text(),
+            type_=sa.String(length=255),
+            existing_nullable=True,
+        )
+        batch_op.alter_column(
+            'refresh_token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+
+    # ### end Alembic commands ###

mage_ai/orchestration/db/models/oauth.py

@@ -12,6 +12,7 @@ from sqlalchemy import (
     ForeignKey,
     Integer,
     String,
+    Text,
     and_,
     asc,
     func,
@@ -324,7 +325,7 @@ class Role(BaseModel):
     @classmethod
     @safe_db_query
     def create_default_roles(
-        self,
+        cls,
         entity: Entity = None,
         entity_id: str = None,
         name_func: Callable[[str], str] = None,
@@ -342,18 +343,18 @@ class Role(BaseModel):
             entity = Entity.GLOBAL
         permissions = Permission.create_default_permissions(entity=entity, entity_id=entity_id)
         mapping = {
-            self.DefaultRole.OWNER: Permission.Access.OWNER,
-            self.DefaultRole.ADMIN: Permission.Access.ADMIN,
-            self.DefaultRole.EDITOR: Permission.Access.EDITOR,
-            self.DefaultRole.VIEWER: Permission.Access.VIEWER,
+            cls.DefaultRole.OWNER: Permission.Access.OWNER,
+            cls.DefaultRole.ADMIN: Permission.Access.ADMIN,
+            cls.DefaultRole.EDITOR: Permission.Access.EDITOR,
+            cls.DefaultRole.VIEWER: Permission.Access.VIEWER,
         }
         for name, access in mapping.items():
             role_name = name
             if name_func is not None:
                 role_name = name_func(name)
-            role = self.query.filter(self.name == role_name).first()
+            role = cls.query.filter(Role.name == role_name).first()
             if not role:
-                self.create(
+                cls.create(
                     name=role_name,
                     permissions=[
                         Permission.query.filter(
@@ -880,10 +881,10 @@ class Oauth2AccessToken(BaseModel):
     expires = Column(DateTime(timezone=True))
     oauth2_application = relationship(Oauth2Application, back_populates='oauth2_access_tokens')
     oauth2_application_id = Column(Integer, ForeignKey('oauth2_application.id'))
-    token = Column(String(255), index=True, unique=True)
+    token = Column(Text, index=True, unique=True)
     user = relationship(User, back_populates='oauth2_access_tokens')
     user_id = Column(Integer, ForeignKey('user.id'))
-    refresh_token = Column(String(255))
+    refresh_token = Column(Text)
 
     def is_valid(self) -> bool:
         return self.token and \

mage_ai/orchestration/db/models/schedules.py

@@ -1913,3 +1913,24 @@ class Backfill(BaseModel):
                 Backfill.pipeline_schedule_id.in_(pipeline_schedule_ids),
             )
         return []
+
+    @property
+    def pipeline_run_status_counts(self) -> Dict:
+        status_counts = dict()
+        execution_dates_counted = set()
+
+        # Sort the pipeline runs by id in reverse order so the first pipeline run
+        # checked is the latest pipeline run created for a given execution date.
+        pipeline_runs_sorted = sorted(
+            self.pipeline_runs,
+            key=lambda pr: (pr.execution_date, pr.id),
+            reverse=True,
+        )
+
+        for pr in pipeline_runs_sorted:
+            # Only count a pipeline run once per execution date
+            if pr.execution_date not in execution_dates_counted:
+                status_counts[pr.status] = (status_counts.get(pr.status) or 0) + 1
+                execution_dates_counted.add(pr.execution_date)
+
+        return status_counts

mage_ai/orchestration/notification/sender.py

@@ -1,4 +1,5 @@
 import os
+import traceback
 from typing import Dict
 
 from mage_ai.orchestration.notification.config import (
@@ -61,33 +62,54 @@ class NotificationSender:
         if summary is None:
             return
         if self.config.slack_config is not None and self.config.slack_config.is_valid:
-            send_slack_message(self.config.slack_config, details or summary, title)
+            try:
+                send_slack_message(self.config.slack_config, details or summary, title)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.teams_config is not None and self.config.teams_config.is_valid:
-            send_teams_message(self.config.teams_config, summary)
+            try:
+                send_teams_message(self.config.teams_config, summary)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.discord_config is not None and self.config.discord_config.is_valid:
-            send_discord_message(self.config.discord_config, summary, title)
+            try:
+                send_discord_message(self.config.discord_config, summary, title)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.telegram_config is not None and self.config.telegram_config.is_valid:
-            send_telegram_message(self.config.telegram_config, summary, title)
+            try:
+                send_telegram_message(self.config.telegram_config, summary, title)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.google_chat_config is not None and self.config.google_chat_config.is_valid:
-            send_google_chat_message(self.config.google_chat_config, summary)
+            try:
+                send_google_chat_message(self.config.google_chat_config, summary)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.email_config is not None and title is not None:
-            send_email(
-                self.config.email_config,
-                subject=title,
-                message=details or summary,
-            )
+            try:
+                send_email(
+                    self.config.email_config,
+                    subject=title,
+                    message=details or summary,
+                )
+            except Exception:
+                traceback.print_exc()
 
         if self.config.opsgenie_config is not None and self.config.opsgenie_config.is_valid:
-            send_opsgenie_alert(
-                self.config.opsgenie_config,
-                message=title,
-                description=details or summary,
-            )
+            try:
+                send_opsgenie_alert(
+                    self.config.opsgenie_config,
+                    message=title,
+                    description=details or summary,
+                )
+            except Exception:
+                traceback.print_exc()
 
     def send_pipeline_run_success_message(self, pipeline, pipeline_run) -> None:
         if AlertOn.PIPELINE_RUN_SUCCESS in self.config.alert_on:

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -284,8 +284,7 @@ class PipelineScheduler:
                     or PipelineRun.PipelineRunStatus.FAILED
                 )
                 self.pipeline_run.update(status=status)
-
-                self.on_pipeline_run_failure('Pipeline run timed out.')
+                self.on_pipeline_run_failure('Pipeline run timed out.', status=status)
             elif self.pipeline_run.any_blocks_failed() and not self.allow_blocks_to_fail:
                 self.pipeline_run.update(
                     status=PipelineRun.PipelineRunStatus.FAILED)
@@ -318,30 +317,38 @@ class PipelineScheduler:
                     self.__schedule_blocks(block_runs)
 
     @safe_db_query
-    def on_pipeline_run_failure(self, error_msg: str) -> None:
-        failed_block_runs = self.pipeline_run.failed_block_runs
-        stacktrace = None
-        for br in failed_block_runs:
-            if br.metrics:
-                message = br.metrics.get('error', {}).get('message')
-                if message:
-                    message_split = message.split('\n')
-                    # Truncate the error message if it has too many lines, set max
-                    # lines at 50
-                    if len(message_split) > 50:
-                        message_split = message_split[-50:]
-                        message_split.insert(0, '... (error truncated)')
-                    message = '\n'.join(message_split)
-                    stacktrace = f'Error for block {br.block_uuid}:\n{message}'
-                    break
-
+    def on_pipeline_run_failure(
+        self,
+        error_msg: str,
+        status=PipelineRun.PipelineRunStatus.FAILED,
+    ) -> None:
         asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
-        self.notification_sender.send_pipeline_run_failure_message(
-            pipeline=self.pipeline,
-            pipeline_run=self.pipeline_run,
-            error=error_msg,
-            stacktrace=stacktrace,
-        )
+
+        if status == PipelineRun.PipelineRunStatus.FAILED:
+            # Only send notification when pipeline run status is FAILED
+            failed_block_runs = self.pipeline_run.failed_block_runs
+            stacktrace = None
+            for br in failed_block_runs:
+                if br.metrics:
+                    message = br.metrics.get('error', {}).get('message')
+                    if message:
+                        message_split = message.split('\n')
+                        # Truncate the error message if it has too many lines, set max
+                        # lines at 50
+                        if len(message_split) > 50:
+                            message_split = message_split[-50:]
+                            message_split.insert(0, '... (error truncated)')
+                        message = '\n'.join(message_split)
+                        stacktrace = f'Error for block {br.block_uuid}:\n{message}'
+                        break
+
+            self.notification_sender.send_pipeline_run_failure_message(
+                pipeline=self.pipeline,
+                pipeline_run=self.pipeline_run,
+                error=error_msg,
+                stacktrace=stacktrace,
+            )
+
         # Cancel block runs that are still in progress for the pipeline run.
         cancel_block_runs_and_jobs(self.pipeline_run, self.pipeline)
 

mage_ai/orchestration/triggers/api.py

@@ -1,7 +1,8 @@
 from datetime import datetime
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Union
 
 from mage_ai.api.resources.PipelineScheduleResource import PipelineScheduleResource
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import ScheduleStatus, ScheduleType
 from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
@@ -25,11 +26,21 @@ def trigger_pipeline(
     poll_timeout: Optional[float] = None,
     schedule_name: str = None,
     verbose: bool = True,
+    remote_blocks: List[Union[Dict, RemoteBlock]] = None,
+    return_remote_blocks: bool = False,
     _should_schedule: bool = False,  # For internal use only (e.g. running hooks from notebook).
 ) -> PipelineRun:
     if variables is None:
         variables = {}
 
+    if remote_blocks:
+        arr = []
+        for remote_block in remote_blocks:
+            if isinstance(remote_block, dict):
+                remote_block = RemoteBlock.load(**remote_block)
+            arr.append(remote_block.to_dict())
+        variables['remote_blocks'] = arr
+
     pipeline = Pipeline.get(pipeline_uuid, all_projects=project_platform_activated())
 
     pipeline_schedule = __fetch_or_create_pipeline_schedule(pipeline, schedule_name=schedule_name)
@@ -50,6 +61,23 @@ def trigger_pipeline(
             verbose=verbose,
         )
 
+    if return_remote_blocks and pipeline_run and pipeline_run.pipeline:
+        pipeline = pipeline_run.pipeline
+
+        return [
+            dict(
+                remote_blocks=[
+                    RemoteBlock.load(
+                        block_uuid=block.uuid,
+                        execution_partition=pipeline_run.execution_partition,
+                        pipeline_uuid=pipeline.uuid,
+                        repo_path=pipeline.repo_path,
+                    )
+                    for block in pipeline.blocks_by_uuid.values()
+                ],
+            ),
+        ]
+
     return pipeline_run
 
 

mage_ai/orchestration/triggers/global_data_product.py

@@ -3,9 +3,10 @@ import json
 from datetime import datetime, timedelta, timezone
 from logging import Logger
 from time import sleep
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from mage_ai.data_preparation.logging.logger import DictLogger
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.global_data_product import GlobalDataProduct
 from mage_ai.data_preparation.models.triggers import ScheduleStatus, ScheduleType
 from mage_ai.orchestration.db import safe_db_query
@@ -145,6 +146,7 @@ def trigger_and_check_status(
     round_number: int = 0,
     verbose: bool = True,
     should_schedule: bool = False,
+    remote_blocks: List[Union[Dict, RemoteBlock]] = None,
 ):
     tags = merge_dict(logging_tags, dict(
         block_uuid=block.uuid if block else None,
@@ -163,6 +165,7 @@ def trigger_and_check_status(
             print(log_message)
             print(json.dumps(tags, indent=2))
 
+    pipeline_run = None
     pipeline_run_created = None
     tries = 0
 
@@ -305,6 +308,7 @@ def trigger_and_check_status(
                         global_data_product.pipeline,
                         pipeline_schedule,
                         dict(variables=variables),
+                        remote_blocks=remote_blocks,
                         should_schedule=should_schedule,
                     )
                     if pipeline_run_created:
@@ -312,11 +316,9 @@ def trigger_and_check_status(
                             f'Created pipeline run {pipeline_run_created.id} for '
                             f'global data product {global_data_product.uuid}.'
                         )
+                        pipeline_run = pipeline_run_created
 
                     lock.release_lock(__lock_key_for_creating_pipeline_run(global_data_product))
-
-                    if pipeline_run_created:
-                        break
             else:
                 __log(
                     f'No pipeline run for global data product {global_data_product.uuid} '
@@ -339,6 +341,7 @@ def trigger_and_check_status(
                     poll_timeout=poll_timeout,
                     verbose=verbose,
                     should_schedule=should_schedule,
+                    remote_blocks=remote_blocks,
                 )
                 break
         else:
@@ -359,6 +362,8 @@ def trigger_and_check_status(
         else:
             break
 
+    return pipeline_run or pipeline_run_created
+
 
 def fetch_or_create_pipeline_schedule(global_data_product: GlobalDataProduct) -> PipelineSchedule:
     pipeline_uuid = global_data_product.object_uuid

mage_ai/orchestration/triggers/utils.py

@@ -1,7 +1,8 @@
 from datetime import datetime, timedelta
 from time import sleep
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Union
 
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import ScheduleStatus
 from mage_ai.orchestration.db import db_connection, safe_db_query
@@ -84,6 +85,7 @@ def create_and_start_pipeline_run(
     pipeline_schedule: PipelineSchedule,
     payload: Dict = None,
     should_schedule: bool = False,
+    remote_blocks: List[Union[Dict, RemoteBlock]] = None,
 ) -> PipelineRun:
     if payload is None:
         payload = {}
@@ -94,6 +96,13 @@ def create_and_start_pipeline_run(
         payload,
     )
 
+    if remote_blocks:
+        variables = configured_payload.get('variables', {})
+        if variables.get('remote_blocks'):
+            remote_blocks = variables.get('remote_blocks', []) + remote_blocks
+        variables['remote_blocks'] = remote_blocks
+        configured_payload['variables'] = variables
+
     pipeline_run = PipelineRun.create(**configured_payload)
 
     # Do not start the pipeline run immediately due to concurrency control