mage-ai 0.9.57__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/command_center/triggers/utils.py

@@ -0,0 +1,463 @@
+import urllib.parse
+from datetime import datetime
+from typing import Dict, List
+from uuid import uuid4
+
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.api.policies.FilePolicy import FilePolicy
+from mage_ai.command_center.constants import (
+    ApplicationType,
+    ButtonActionType,
+    ItemType,
+    ObjectType,
+)
+from mage_ai.command_center.models import TriggerMetadata
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.models.triggers import (
+    SCHEDULE_TYPE_TO_LABEL,
+    ScheduleInterval,
+    ScheduleStatus,
+    ScheduleType,
+)
+from mage_ai.orchestration.db.models.schedules import PipelineSchedule
+from mage_ai.presenters.interactions.constants import InteractionInputType
+
+
+def add_application_actions(item_dict: Dict) -> Dict:
+    model = item_dict['metadata']['trigger']
+    model_id = model['id']
+    pipeline_uuid = model['pipeline_uuid']
+
+    application_action = dict(
+        request=dict(
+            operation=OperationType.DETAIL,
+            resource='pipeline_schedules',
+            resource_id=model_id,
+            response_resource_key='pipeline_schedule',
+            resource_parent='pipelines',
+            resource_parent_id=urllib.parse.quote_plus(pipeline_uuid),
+            query=dict(
+                _format='with_runtime_average',
+            ),
+        ),
+        uuid='model_detail',
+    )
+
+    return dict(
+        applications=[
+            dict(
+                application_type=ApplicationType.DETAIL_LIST,
+                actions=[application_action],
+                uuid='model_detail_list',
+            ),
+            dict(
+                application_type=ApplicationType.DETAIL,
+                actions=[application_action],
+                buttons=[
+                    dict(
+                        label='Open trigger',
+                        keyboard_shortcuts=[[13]],
+                        action_types=[
+                            ButtonActionType.CUSTOM_ACTIONS,
+                        ],
+                        actions=[
+                            dict(
+                                page=dict(
+                                    path=f'/pipelines/{pipeline_uuid}/triggers/{model_id}',
+                                ),
+                                uuid='open_model',
+                            ),
+                        ],
+                    ),
+                ],
+                uuid='model_detail',
+            ),
+        ],
+    )
+
+
+async def build_and_score(
+    factory,
+    model: PipelineSchedule,
+    items: List[Dict],
+    add_application: bool = False,
+):
+    item_dict = dict(
+        item_type=ItemType.DETAIL,
+        object_type=ObjectType.TRIGGER,
+        title=model.name,
+        description=model.description or ' '.join([s for s in [
+            SCHEDULE_TYPE_TO_LABEL.get(model.schedule_type),
+            model.schedule_interval,
+            model.status,
+        ] if s]),
+        uuid=f'{model.id}-{model.pipeline_uuid}-{model.repo_path}',
+        metadata=dict(
+            trigger=dict(
+                description=model.description,
+                global_data_product_uuid=model.global_data_product_uuid,
+                id=model.id,
+                name=model.name,
+                pipeline_uuid=model.pipeline_uuid,
+                repo_path=model.repo_path,
+                schedule_interval=model.schedule_interval,
+                schedule_type=model.schedule_type,
+                settings=model.settings,
+                sla=model.sla,
+                start_time=model.start_time,
+                status=model.status,
+            ),
+        ),
+        display_settings_by_attribute=dict(
+            description=dict(
+                text_styles=dict(
+                    monospace=True,
+                    small=True,
+                ),
+            ),
+        ),
+    )
+
+    if add_application:
+        item_dict.update(add_application_actions(item_dict))
+
+    scored = factory.filter_score(item_dict)
+    if scored:
+        items.append(scored)
+
+
+async def build_create_and_score(factory, model: Pipeline) -> Dict:
+    item_dict = dict(
+        item_type=ItemType.CREATE,
+        object_type=ObjectType.TRIGGER,
+        title='Create new trigger',
+        description=f'Create a new trigger for pipeline {model.name or model.uuid}',
+        uuid=f'{model.uuid}_{ItemType.CREATE}_{ObjectType.TRIGGER}',
+        metadata=dict(
+            pipeline=dict(
+                repo_path=model.repo_path,
+                uuid=model.uuid,
+            ),
+        ),
+        display_settings_by_attribute=dict(
+            icon=dict(
+                color_uuid='accent.negative',
+                icon_uuid='Lightning',
+            ),
+        ),
+        applications=[
+            dict(
+                application_type=ApplicationType.FORM,
+                buttons=[
+                    dict(
+                        label='Cancel',
+                        tooltip='Discard changes and go back.',
+                        keyboard_shortcuts=[['metaKey', 27]],
+                        action_types=[
+                            ButtonActionType.RESET_FORM,
+                            ButtonActionType.CLOSE_APPLICATION,
+                        ],
+                    ),
+                    dict(
+                        label='Create trigger',
+                        tooltip='Save changes and create the new trigger.',
+                        keyboard_shortcuts=[[13]],
+                        action_types=[
+                            ButtonActionType.EXECUTE,
+                            ButtonActionType.RESET_FORM,
+                        ],
+                    ),
+                ],
+                settings=[
+                    dict(
+                        label='Name',
+                        placeholder='e.g. Daily ETL',
+                        display_settings=dict(
+                            icon_uuid='Alphabet',
+                        ),
+                        name='request.payload.pipeline_schedule.name',
+                        type=InteractionInputType.TEXT_FIELD,
+                        required=True,
+                        action_uuid='create_model',
+                    ),
+                    dict(
+                        label='Description',
+                        placeholder='e.g. Runs daily for global users.',
+                        display_settings=dict(
+                            icon_uuid='Alphabet',
+                        ),
+                        name='request.payload.pipeline_schedule.description',
+                        type=InteractionInputType.TEXT_FIELD,
+                        required=True,
+                        action_uuid='create_model',
+                    ),
+                    dict(
+                        label='Type',
+                        description='The type of trigger to create.',
+                        placeholder='e.g. Schedule',
+                        name='request.payload.pipeline_schedule.schedule_type',
+                        type=InteractionInputType.DROPDOWN_MENU,
+                        options=[dict(
+                            label=SCHEDULE_TYPE_TO_LABEL[schedule_type],
+                            value=schedule_type,
+                        ) for schedule_type in ScheduleType if (
+                            schedule_type in SCHEDULE_TYPE_TO_LABEL
+                        )],
+                        value=ScheduleType.TIME.value,
+                        required=True,
+                        action_uuid='create_model',
+                    ),
+                ],
+                uuid='model_create',
+            ),
+        ],
+        actions=[
+            dict(
+                request=dict(
+                    operation=OperationType.CREATE,
+                    payload=dict(
+                        pipeline_schedule=dict(
+                            description=None,
+                            name=None,
+                            schedule_type=None,
+                        ),
+                    ),
+                    resource='pipeline_schedules',
+                    response_resource_key='pipeline_schedule',
+                    resource_parent='pipelines',
+                    resource_parent_id=urllib.parse.quote_plus(model.uuid),
+                ),
+                uuid='create_model',
+            ),
+            dict(
+                upstream_action_value_key_mapping=dict(
+                    create_model={
+                        'data.pipeline_schedule.id': 'page.parameters.pipeline_schedule_id',
+                    }
+                ),
+                page=dict(
+                    path=(
+                        f'/pipelines/{urllib.parse.quote_plus(model.uuid)}'
+                        '/triggers/:pipeline_schedule_id/edit'
+                    ),
+                    parameters=dict(
+                        pipeline_schedule_id=None,
+                    ),
+                ),
+                uuid='open_model',
+            ),
+        ],
+        condition=lambda opts: FilePolicy(
+            None,
+            opts.get('user'),
+        ).has_at_least_editor_role(),
+    )
+
+    scored = factory.filter_score(item_dict)
+    if scored:
+        return
+
+
+async def build_run_once_and_score(factory, model: Pipeline) -> Dict:
+    item_dict = dict(
+        item_type=ItemType.CREATE,
+        object_type=ObjectType.TRIGGER,
+        title='Trigger pipeline once',
+        description=f'Run pipeline {model.name or model.uuid} once',
+        uuid=f'{model.uuid}_{ItemType.CREATE}_{ObjectType.TRIGGER}',
+        metadata=dict(
+            pipeline=dict(
+                repo_path=model.repo_path,
+                uuid=model.uuid,
+            ),
+        ),
+        display_settings_by_attribute=dict(
+            icon=dict(
+                color_uuid='accent.negative',
+                icon_uuid='Lightning',
+            ),
+        ),
+        applications=[
+            dict(
+                application_type=ApplicationType.FORM,
+                buttons=[
+                    dict(
+                        label='Cancel',
+                        tooltip='Discard changes and go back.',
+                        keyboard_shortcuts=[['metaKey', 27]],
+                        action_types=[
+                            ButtonActionType.RESET_FORM,
+                            ButtonActionType.CLOSE_APPLICATION,
+                        ],
+                    ),
+                    dict(
+                        label='Trigger pipeline',
+                        tooltip='Run pipeline now.',
+                        action_types=[
+                            ButtonActionType.EXECUTE,
+                            ButtonActionType.RESET_FORM,
+                        ],
+                    ),
+                ],
+                settings=[
+                    dict(
+                        label='Variables (optional)',
+                        description='JSON containing variables for this pipeline run.',
+                        name='request.payload.pipeline_schedule.variables',
+                        type=InteractionInputType.CODE,
+                        action_uuid='create_model',
+                        style=dict(
+                            language='json',
+                        ),
+                    ),
+                ],
+                uuid='model_create',
+            ),
+        ],
+        actions=[
+            dict(
+                request=dict(
+                    operation=OperationType.CREATE,
+                    payload=dict(
+                        pipeline_schedule=dict(
+                            name=f'Run once {uuid4().hex[:5]}',
+                            schedule_interval=ScheduleInterval.ONCE,
+                            schedule_type=ScheduleType.TIME,
+                            start_time=datetime.now().isoformat(),
+                            status=ScheduleStatus.ACTIVE,
+                            variables=None,
+                        ),
+                    ),
+                    resource='pipeline_schedules',
+                    response_resource_key='pipeline_schedule',
+                    resource_parent='pipelines',
+                    resource_parent_id=urllib.parse.quote_plus(model.uuid),
+                ),
+                uuid='create_model',
+            ),
+            dict(
+                upstream_action_value_key_mapping=dict(
+                    create_model={
+                        'data.pipeline_schedule.id': 'page.parameters.pipeline_schedule_id',
+                    }
+                ),
+                page=dict(
+                    path=(
+                        f'/pipelines/{urllib.parse.quote_plus(model.uuid)}'
+                        '/triggers/:pipeline_schedule_id'
+                    ),
+                    parameters=dict(
+                        pipeline_schedule_id=None,
+                    ),
+                ),
+                uuid='open_model',
+            ),
+        ],
+        condition=lambda opts: FilePolicy(
+            None,
+            opts.get('user'),
+        ).has_at_least_editor_role(),
+    )
+
+    scored = factory.filter_score(item_dict)
+    return scored
+
+
+async def build_create_pipeline_run_once(factory, model: TriggerMetadata) -> Dict:
+    item_dict = dict(
+        item_type=ItemType.CREATE,
+        object_type=ObjectType.PIPELINE_RUN,
+        title='Run trigger once',
+        description=f'Create a 1 time run for trigger {model.name or model.id}',
+        uuid=f'{model.id}_{ItemType.CREATE}_{ObjectType.PIPELINE_RUN}',
+        metadata=dict(
+            trigger=model.to_dict(),
+        ),
+        display_settings_by_attribute=dict(
+            icon=dict(
+                color_uuid='accent.negative',
+                icon_uuid='Lightning',
+            ),
+        ),
+        applications=[
+            dict(
+                application_type=ApplicationType.FORM,
+                buttons=[
+                    dict(
+                        label='Cancel',
+                        tooltip='Discard changes and go back.',
+                        keyboard_shortcuts=[['metaKey', 27]],
+                        action_types=[
+                            ButtonActionType.RESET_FORM,
+                            ButtonActionType.CLOSE_APPLICATION,
+                        ],
+                    ),
+                    dict(
+                        label='Run now',
+                        tooltip='Run this trigger now.',
+                        keyboard_shortcuts=[[18, 13]],
+                        action_types=[
+                            ButtonActionType.EXECUTE,
+                            ButtonActionType.RESET_FORM,
+                        ],
+                    ),
+                ],
+                settings=[
+                    dict(
+                        label='Variables (optional)',
+                        description='JSON containing variables for this pipeline run.',
+                        name='request.payload.pipeline_run.variables',
+                        type=InteractionInputType.CODE,
+                        action_uuid='create_model',
+                        style=dict(
+                            language='json',
+                        ),
+                    ),
+                ],
+                uuid='model_create',
+            ),
+        ],
+        actions=[
+            dict(
+                request=dict(
+                    operation=OperationType.CREATE,
+                    payload=dict(
+                        pipeline_run=dict(
+                            pipeline_schedule_id=model.id,
+                            pipeline_uuid=model.pipeline_uuid,
+                            variables=None,
+                        ),
+                    ),
+                    resource='pipeline_runs',
+                    response_resource_key='pipeline_run',
+                    resource_parent='pipeline_schedules',
+                    resource_parent_id=model.id,
+                ),
+                uuid='create_model',
+            ),
+            dict(
+                upstream_action_value_key_mapping=dict(
+                    create_model={
+                        'data.pipeline_run.id': 'page.parameters.pipeline_run_id',
+                    }
+                ),
+                page=dict(
+                    path=(
+                        f'/pipelines/{urllib.parse.quote_plus(model.pipeline_uuid)}'
+                        '/runs/:pipeline_run_id'
+                    ),
+                    parameters=dict(
+                        pipeline_run_id=None,
+                    ),
+                ),
+                uuid='open_model',
+            ),
+        ],
+        condition=lambda opts: FilePolicy(
+            None,
+            opts.get('user'),
+        ).has_at_least_editor_role(),
+    )
+
+    scored = factory.filter_score(item_dict)
+    return scored

mage_ai/command_center/utils.py

@@ -0,0 +1,30 @@
+import os
+from pathlib import Path
+from typing import Dict
+
+
+def shorten_directory(full_path: str) -> Dict:
+    parts = Path(os.path.dirname(full_path)).parts
+    parts_count = len(parts)
+    dir_name = ''
+    if parts_count >= 1:
+        os.path.join(*parts[max([0, parts_count - 3]):])
+
+    if parts_count >= 4:
+        if parts_count >= 5:
+            dir_name = os.path.join('.', dir_name)
+        else:
+            dir_name = os.path.join('..', dir_name)
+
+        for _i in range(min(1, parts_count - 4)):
+            dir_name = os.path.join('..', dir_name)
+
+    extension = Path(full_path or '').suffix
+    if extension:
+        extension = extension.replace('.', '')
+
+    return dict(
+        directory=dir_name,
+        extension=extension,
+        parts=parts,
+    )

mage_ai/data_preparation/executors/azure_container_instance_executor.py

@@ -8,7 +8,7 @@ from mage_ai.shared.hash import merge_dict
 class AzureContainerInstanceExecutor(BlockExecutor):
     RETRYABLE = False
 
-    def __init__(self, pipeline, block_uuid: str, execution_partition: str = None):
+    def __init__(self, pipeline, block_uuid: str, execution_partition: str = None, **kwargs):
         super().__init__(pipeline, block_uuid, execution_partition=execution_partition)
         self.executor_config = self.pipeline.repo_config.azure_container_instance_config or dict()
         if self.block.executor_config is not None:

mage_ai/data_preparation/executors/block_executor.py

@@ -23,15 +23,15 @@ from mage_ai.data_preparation.models.block.data_integration.utils import (
 from mage_ai.data_preparation.models.block.dynamic.dynamic_child import (
     DynamicChildBlockFactory,
 )
-from mage_ai.data_preparation.models.block.utils import (
-    dynamic_block_uuid as dynamic_block_uuid_func,
-)
-from mage_ai.data_preparation.models.block.utils import (
-    dynamic_block_values_and_metadata,
+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.block.utils import (
+    build_dynamic_block_uuid,
+    dynamic_block_values_and_metadata,
+)
 from mage_ai.data_preparation.models.constants import (
     BlockLanguage,
     BlockType,
@@ -42,6 +42,7 @@ from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.data_preparation.models.triggers import ScheduleInterval, ScheduleType
 from mage_ai.data_preparation.shared.retry import RetryConfig
 from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun
+from mage_ai.shared.custom_logger import DX_PRINTER
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.utils import clean_name
 
@@ -57,7 +58,8 @@ class BlockExecutor:
         self,
         pipeline,
         block_uuid,
-        execution_partition=None
+        execution_partition: str = None,
+        block_run_id: int = None,
     ):
         """
         Initialize the BlockExecutor.
@@ -75,18 +77,53 @@ class BlockExecutor:
             block_uuid=clean_name(self.block_uuid),
             partition=self.execution_partition,
             repo_config=self.pipeline.repo_config,
+            repo_path=self.pipeline.repo_path,
         )
         self.logger = DictLogger(self.logger_manager.logger)
         self.project = Project(self.pipeline.repo_config)
+        self.retry_metadata = dict(attempts=0)
 
         self.block = self.pipeline.get_block(self.block_uuid, check_template=True)
 
         # If this is the original block run for the original dynamic block
-        if self.block and \
-                self.block.uuid == self.block_uuid and \
-                is_dynamic_block_child(self.block):
 
-            self.block = DynamicChildBlockFactory(self.block)
+        # Check to see if this block is the original dynamic child block or
+        # a clone of the original dynamic child block.
+        factory = DynamicChildBlockFactory(self.block, block_run_id=block_run_id)
+        wrapper = factory.wrapper()
+        if self.block:
+            is_clone_of_original = wrapper.is_clone_of_original()
+            is_dynamic = wrapper.is_dynamic()
+            is_dynamic_child = wrapper.is_dynamic_child()
+            is_original = wrapper.is_original()
+            is_replicated = wrapper.is_replicated()
+
+            DX_PRINTER.info(
+                (
+                    f'Checking if block run {block_uuid} is dynamic child and '
+                    'original, clone of original, or replicated'
+                ),
+                block=self.block,
+                is_clone_of_original=is_clone_of_original,
+                is_dynamic=is_dynamic,
+                is_dynamic_child=is_dynamic_child,
+                is_original=is_original,
+                is_replicated=is_replicated,
+                __uuid='BlockExecutor',
+            )
+
+            if is_dynamic_child and (is_original or is_clone_of_original):
+                self.block = factory
+
+                DX_PRINTER.info(
+                    'Initializing dynamic child block factory',
+                    block=self.block,
+                    clone_of_original=wrapper.is_clone_of_original(),
+                    is_dynamic=wrapper.is_dynamic(),
+                    is_dynamic_child=wrapper.is_dynamic_child(),
+                    original=wrapper.is_original(),
+                    __uuid='BlockExecutor',
+                )
 
         self.block_run = None
 
@@ -327,8 +364,8 @@ class BlockExecutor:
 
                         values, block_metadata = dynamic_block_values_and_metadata(
                             block_grandparent,
-                            self.execution_partition,
-                            block_grandparent_uuid,
+                            block_uuid=block_grandparent_uuid,
+                            execution_partition=self.execution_partition,
                         )
 
                         for idx, _ in enumerate(values):
@@ -338,7 +375,7 @@ class BlockExecutor:
                                 metadata = {}
 
                             dynamic_upstream_block_uuids_reduce.append(
-                                dynamic_block_uuid_func(
+                                build_dynamic_block_uuid(
                                     upstream_block.uuid,
                                     metadata,
                                     idx,
@@ -555,6 +592,7 @@ class BlockExecutor:
                         exponential_backoff=retry_config.exponential_backoff,
                         logger=self.logger,
                         logging_tags=tags,
+                        retry_metadata=self.retry_metadata,
                     )
                     def __execute_with_retry():
                         return self._execute(
@@ -563,7 +601,9 @@ class BlockExecutor:
                             block_run_outputs_cache=block_run_outputs_cache,
                             cache_block_output_in_memory=cache_block_output_in_memory,
                             callback_url=callback_url,
-                            global_vars=global_vars,
+                            global_vars=merge_dict(global_vars or {}, dict(
+                                retry=self.retry_metadata,
+                            )),
                             update_status=update_status,
                             input_from_output=input_from_output,
                             logging_tags=tags,
@@ -614,7 +654,7 @@ class BlockExecutor:
                         )
                     self._execute_callback(
                         'on_failure',
-                        callback_kwargs=dict(__error=error),
+                        callback_kwargs=dict(__error=error, retry=self.retry_metadata),
                         dynamic_block_index=dynamic_block_index,
                         dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                         global_vars=global_vars,
@@ -644,8 +684,12 @@ class BlockExecutor:
                 if on_complete is not None:
                     metrics = None
                     if self.block and is_dynamic_block(self.block):
-                        if result and result.get('output'):
-                            output = result.get('output')
+                        if result:
+                            if isinstance(result, dict):
+                                output = result.get('output')
+                            else:
+                                output = result
+
                             if len(output) >= 1:
                                 child_data = output[0]
                                 metrics = dict(children=len(child_data))
@@ -670,6 +714,7 @@ class BlockExecutor:
             if not data_integration_metadata or is_original_block:
                 self._execute_callback(
                     'on_success',
+                    callback_kwargs=dict(retry=self.retry_metadata),
                     dynamic_block_index=dynamic_block_index,
                     dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                     global_vars=global_vars,
@@ -731,6 +776,10 @@ class BlockExecutor:
             logging_tags = dict()
 
         extra_options = {}
+
+        if self.block_run and self.block_run.metrics:
+            extra_options['metadata'] = self.block_run.metrics.get('metadata')
+
         if cache_block_output_in_memory:
             store_variables = False
         else:
@@ -1030,8 +1079,8 @@ class BlockExecutor:
 
                     return arr
 
-        if is_dynamic_block_child(self.block) and self.block_uuid == self.block.uuid:
-            self.block.pipeline_run = pipeline_run
+        if self.block and isinstance(self.block, DynamicChildBlockFactory):
+            self.block.set_pipeline_run(pipeline_run)
 
         result = self.block.execute_sync(
             analyze_outputs=analyze_outputs,