mage-ai 0.9.45__py3-none-any.whl → 0.9.47__py3-none-any.whl

mage_ai/data_preparation/models/global_hooks/predicates.py

@@ -0,0 +1,316 @@
+import json
+from abc import abstractmethod
+from collections.abc import Iterable
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Union
+
+from mage_ai.api.resources.BaseResource import BaseResource
+from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.global_hooks.constants import (
+    INTERNAL_DEFAULT_PREDICATE_VALUE,
+    PredicateAndOrOperator,
+    PredicateObjectType,
+    PredicateOperator,
+    PredicateValueDataType,
+)
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class PredicateValueType(BaseDataClass):
+    value_data_type: PredicateValueDataType = None
+    value_type: Dict = None
+
+    def __post_init__(self):
+        self.serialize_attribute_enum('value_data_type', PredicateValueDataType)
+
+        if self.value_type and isinstance(self.value_type, dict):
+            self.value_type = self.__class__.load(**self.value_type)
+
+    def to_dict(self, **kwargs) -> Dict:
+        data = super().to_dict(**kwargs)
+
+        if self.value_type and issubclass(self.value_type.__class__, BaseDataClass):
+            data['value_type'] = self.value_type.to_dict(**kwargs)
+
+        return data
+
+
+@dataclass
+class BasePredicate(BaseDataClass):
+    and_or_operator: PredicateAndOrOperator = None
+    left_object_keys: List[str] = None
+    left_object_type: PredicateObjectType = None
+    left_value: Union[bool, Dict, float, int, List, str] = None
+    left_value_type: PredicateValueType = None
+    operator: PredicateOperator = None
+    predicates: List = None
+    right_object_keys: List[str] = None
+    right_object_type: PredicateObjectType = None
+    right_value: Union[bool, Dict, float, int, List, str] = None
+    right_value_type: PredicateValueType = None
+
+    @abstractmethod
+    def validate(
+        self,
+        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        error: Dict = None,  # code, errors, message, type
+        hook: Dict = None,
+        meta: Dict = None,
+        metadata: Dict = None,
+        payload: Dict = None,
+        query: Dict = None,
+        resource: Dict = None,
+        resource_id: Union[int, str] = None,
+        resource_parent_id: Union[int, str] = None,
+        resources: Dict = None,
+        user: Dict = None,
+    ) -> bool:
+        return False
+
+
+@dataclass
+class HookPredicate(BasePredicate):
+    predicates: List[BasePredicate] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_classes('predicates', self.__class__)
+        self.serialize_attribute_class('left_value_type', PredicateValueType)
+        self.serialize_attribute_class('right_value_type', PredicateValueType)
+        self.serialize_attribute_enum('and_or_operator', PredicateAndOrOperator)
+        self.serialize_attribute_enum('left_object_type', PredicateObjectType)
+        self.serialize_attribute_enum('operator', PredicateOperator)
+        self.serialize_attribute_enum('right_object_type', PredicateObjectType)
+
+    def validate(
+        self,
+        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        **kwargs,
+    ) -> bool:
+        if self.predicates:
+            operator_func = all
+
+            if PredicateAndOrOperator.OR == self.and_or_operator:
+                operator_func = any
+
+            return operator_func([predicate.validate(
+                operation_resource,
+                **kwargs,
+            ) for predicate in self.predicates])
+
+        return self.__validate_resource(operation_resource, **kwargs)
+
+    def __validate_resource(
+        self,
+        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        **kwargs,
+    ) -> bool:
+        if isinstance(operation_resource, Iterable) and not isinstance(operation_resource, dict):
+            return all([self.__validate(
+                operation_resource=res,
+                **kwargs,
+            ) for res in operation_resource])
+
+        return self.__validate(operation_resource=operation_resource, **kwargs)
+
+    def __validate(self, **kwargs) -> bool:
+
+        left_object = None
+        right_object = None
+
+        left_value_to_compare = INTERNAL_DEFAULT_PREDICATE_VALUE
+        right_value_to_compare = INTERNAL_DEFAULT_PREDICATE_VALUE
+
+        # Get the value from the object
+        if self.left_object_type:
+            key = self.left_object_type.value
+            if key in kwargs:
+                left_object = kwargs[key]
+
+        if self.right_object_type:
+            key = self.right_object_type.value
+            if key in kwargs:
+                right_object = kwargs[key]
+
+        # Get value from object using keys if self.right_value is not None
+        if self.left_value is not None:
+            left_value_to_compare = self.left_value
+        elif left_object:
+            if self.left_object_keys:
+                left_value_to_compare = get_value(left_object, self.left_object_keys)
+            else:
+                left_value_to_compare = left_object
+
+        # Convert left value to the specified data type
+        if self.left_value_type and INTERNAL_DEFAULT_PREDICATE_VALUE != left_value_to_compare:
+            left_value_to_compare = convert_value(left_value_to_compare, self.left_value_type)
+
+        # Get value from object using keys if self.right_value is not None
+        if self.right_value is not None:
+            right_value_to_compare = self.right_value
+        elif right_object:
+            if self.right_object_keys:
+                right_value_to_compare = get_value(right_object, self.right_object_keys)
+            else:
+                right_value_to_compare = right_object
+
+        # Convert value to the specified data type
+        if self.right_value_type and INTERNAL_DEFAULT_PREDICATE_VALUE != right_value_to_compare:
+            right_value_to_compare = convert_value(right_value_to_compare, self.right_value_type)
+
+        if INTERNAL_DEFAULT_PREDICATE_VALUE == left_value_to_compare:
+            left_value_to_compare = None
+
+        if INTERNAL_DEFAULT_PREDICATE_VALUE == right_value_to_compare:
+            right_value_to_compare = None
+
+        if self.operator:
+            return compare_using_operator(
+                self.operator,
+                left_value_to_compare,
+                right_value_to_compare,
+            )
+
+        return False
+
+
+def get_value(
+    object_arg: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline, int, str],
+    keys: List[str] = None,
+) -> Union[bool, Dict, float, int, List, str]:
+    value_temp = None
+
+    if keys:
+        if not isinstance(object_arg, Iterable) and (
+            issubclass(object_arg.__class__, BaseResource) or
+            isinstance(object_arg, Block) or
+            isinstance(object_arg, Pipeline)
+        ):
+            for idx, key in enumerate(keys):
+                if idx == 0:
+                    value_temp = getattr(object_arg, key)
+                else:
+                    value_temp = get_value(value_temp, [key])
+        elif isinstance(object_arg, Iterable) and \
+                not isinstance(object_arg, dict) and \
+                not isinstance(object_arg, str):
+
+            for idx, key in enumerate(keys):
+                if idx == 0:
+                    if isinstance(key, str) and hasattr(object_arg, key):
+                        value_temp = getattr(object_arg, key)
+                    elif isinstance(key, int) or isinstance(key, float):
+                        index_inner = int(key)
+
+                        if object_arg and index_inner < len(object_arg):
+                            value_temp = object_arg[index_inner]
+                else:
+                    value_temp = get_value(value_temp, [key])
+        elif isinstance(object_arg, dict):
+            for idx, key in enumerate(keys):
+                if idx == 0:
+                    value_temp = object_arg.get(key)
+                else:
+                    value_temp = get_value(value_temp, [key])
+
+    if callable(value_temp):
+        value_temp = value_temp()
+
+    return value_temp
+
+
+def convert_value(
+    value: Union[bool, Dict, float, int, List, str],
+    value_type: PredicateValueType,
+) -> Union[bool, Dict, float, int, List, str]:
+    value_data_type = value_type.value_data_type
+    value_type_sub = value_type.value_type
+
+    if PredicateValueDataType.BOOLEAN == value_data_type:
+        value = bool(value)
+    elif PredicateValueDataType.DICTIONARY == value_data_type:
+        if value is not None:
+            if isinstance(value, str):
+                value = json.loads(value)
+            if value_type_sub:
+                for k, v in value.items():
+                    value[k] = convert_value(v, value_type_sub)
+    elif PredicateValueDataType.FLOAT == value_data_type:
+        if value is not None:
+            value = float(value)
+    elif PredicateValueDataType.INTEGER == value_data_type:
+        if value is not None:
+            value = int(value)
+    elif PredicateValueDataType.LIST == value_data_type:
+        if value is not None:
+            if isinstance(value, str):
+                value = json.loads(value)
+            if value_type_sub:
+                value = [convert_value(v, value_type_sub) for v in value]
+    elif PredicateValueDataType.STRING == value_data_type:
+        if value is None:
+            value = ''
+        if isinstance(value, Enum):
+            value = value.value
+        value = str(value)
+
+    return value
+
+
+def compare_using_operator(
+    operator: PredicateOperator,
+    left_value: Union[bool, Dict, float, int, List, str],
+    right_value: Union[bool, Dict, float, int, List, str],
+) -> bool:
+    try:
+        if operator in [
+            PredicateOperator.GREATER_THAN,
+            PredicateOperator.GREATER_THAN_OR_EQUALS,
+            PredicateOperator.LESS_THAN,
+            PredicateOperator.LESS_THAN_OR_EQUALS,
+        ]:
+            if (
+                left_value is not None and
+                not isinstance(left_value, str) and
+                (isinstance(left_value, dict) or isinstance(left_value, Iterable))
+            ) or (
+                right_value is not None and
+                not isinstance(right_value, str) and
+                (isinstance(right_value, dict) or isinstance(right_value, Iterable))
+            ):
+                if left_value is None:
+                    left_value = 0
+                else:
+                    left_value = len(left_value)
+
+                if right_value is None:
+                    right_value = 0
+                else:
+                    right_value = len(right_value)
+
+        if PredicateOperator.EQUALS == operator:
+            return left_value == right_value
+        if PredicateOperator.GREATER_THAN == operator:
+            return left_value > right_value
+        if PredicateOperator.GREATER_THAN_OR_EQUALS == operator:
+            return left_value >= right_value
+        if PredicateOperator.INCLUDES == operator:
+            return left_value in right_value
+        if PredicateOperator.LESS_THAN == operator:
+            return left_value < right_value
+        if PredicateOperator.LESS_THAN_OR_EQUALS == operator:
+            return left_value <= right_value
+        if PredicateOperator.NOT_EQUALS == operator:
+            return left_value != right_value
+        if PredicateOperator.NOT_INCLUDES == operator:
+            return left_value not in right_value
+        if PredicateOperator.NOT_PRESENT == operator:
+            return False if left_value else True
+        if PredicateOperator.PRESENT == operator:
+            return True if left_value else False
+    except TypeError as err:
+        print(f'[ERROR] Predicate error comparing {left_value} {operator} {right_value}: {err}.')
+
+    return False

mage_ai/data_preparation/models/global_hooks/utils.py

@@ -0,0 +1,21 @@
+from typing import Dict
+
+from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.data_preparation.models.global_hooks.constants import (
+    DISABLED_RESOURCE_TYPES,
+    RESTRICTED_RESOURCE_TYPES,
+    VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL,
+    VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED,
+)
+from mage_ai.shared.hash import extract
+
+
+def extract_valid_data(resource_type: EntityName, input_data: Dict) -> Dict:
+    types = DISABLED_RESOURCE_TYPES + RESTRICTED_RESOURCE_TYPES
+    if not resource_type or resource_type in types:
+        return extract(input_data or {}, VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED)
+
+    return extract(
+        input_data or {},
+        VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL,
+    )

mage_ai/data_preparation/models/pipeline.py

@@ -6,10 +6,12 @@ import shutil
 from typing import Any, Callable, Dict, List
 
 import aiofiles
+import dateutil.parser
 import pytz
 import yaml
 from jinja2 import Template
 
+from mage_ai.authentication.permissions.constants import EntityName
 from mage_ai.data_preparation.models.block import Block, run_blocks, run_blocks_sync
 from mage_ai.data_preparation.models.block.data_integration.utils import (
     convert_outputs_to_data,
@@ -30,6 +32,9 @@ from mage_ai.data_preparation.models.constants import (
 )
 from mage_ai.data_preparation.models.errors import SerializationError
 from mage_ai.data_preparation.models.file import File
+from mage_ai.data_preparation.models.pipelines.models import PipelineSettings
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.data_preparation.models.utils import is_yaml_serializable
 from mage_ai.data_preparation.models.variable import Variable
 from mage_ai.data_preparation.repo_manager import (
@@ -72,22 +77,21 @@ class Pipeline:
         self.retry_config = {}
         self.run_pipeline_in_one_process = False
         self.schedules = []
+        self.settings = {}
         self.tags = []
         self.type = PipelineType.PYTHON
-        self.updated_at = datetime.datetime.now(tz=pytz.UTC)
+        self.updated_at = datetime.datetime.utcnow().replace(tzinfo=pytz.UTC)
         self.uuid = uuid
         self.widget_configs = []
         self._executor_count = 1  # Used by streaming pipeline to launch multiple executors
-        if config is None:
-            self.load_config_from_yaml()
-        else:
-            self.load_config(config, catalog=catalog)
+
         if repo_config is None:
             self.repo_config = get_repo_config(repo_path=self.repo_path)
         elif type(repo_config) is dict:
             self.repo_config = RepoConfig.from_dict(repo_config)
         else:
             self.repo_config = repo_config
+
         self.variable_manager = VariableManager.get_manager(
             self.repo_path,
             self.remote_variables_dir or self.variables_dir,
@@ -96,6 +100,11 @@ class Pipeline:
         # Used for showing the operation history. For example: recently viewed pipelines.
         self.history = []
 
+        if config is None:
+            self.load_config_from_yaml()
+        else:
+            self.load_config(config, catalog=catalog)
+
     @property
     def config_path(self):
         return os.path.join(
@@ -421,6 +430,7 @@ class Pipeline:
         global_vars=None,
         run_sensors: bool = True,
         run_tests: bool = True,
+        update_status: bool = True,
     ) -> None:
         """
         Function for synchronous block processing.
@@ -454,6 +464,7 @@ class Pipeline:
                 global_vars=global_vars,
                 run_sensors=run_sensors,
                 run_tests=run_tests,
+                update_status=update_status,
             )
 
     def get_config_from_yaml(self):
@@ -491,6 +502,8 @@ class Pipeline:
             pass
         self.created_at = config.get('created_at')
         self.updated_at = config.get('updated_at')
+        if self.updated_at and isinstance(self.updated_at, str):
+            self.updated_at = dateutil.parser.parse(self.updated_at).replace(tzinfo=pytz.UTC)
         self.type = config.get('type') or self.type
 
         self.block_configs = config.get('blocks') or []
@@ -502,6 +515,7 @@ class Pipeline:
         self.notification_config = config.get('notification_config') or {}
         self.retry_config = config.get('retry_config') or {}
         self.run_pipeline_in_one_process = config.get('run_pipeline_in_one_process', False)
+        self.settings = PipelineSettings.load(**config.get('settings') or {})
         self.spark_config = config.get('spark_config') or {}
         self.tags = config.get('tags') or []
         self.widget_configs = config.get('widgets') or []
@@ -525,6 +539,7 @@ class Pipeline:
                 language=c.get('language'),
                 pipeline=self,
                 replicated_block=c.get('replicated_block'),
+                repo_config=self.repo_config,
                 retry_config=c.get('retry_config'),
                 status=c.get('status'),
                 timeout=c.get('timeout'),
@@ -614,6 +629,10 @@ class Pipeline:
         return blocks_by_uuid
 
     def to_dict_base(self, exclude_data_integration=False) -> Dict:
+        updated_at = self.updated_at
+        if updated_at and hasattr(updated_at, 'isoformat'):
+            updated_at = updated_at.isoformat()
+
         base = dict(
             concurrency_config=self.concurrency_config,
             created_at=self.created_at,
@@ -624,12 +643,15 @@ class Pipeline:
             executor_type=self.executor_type,
             name=self.name,
             notification_config=self.notification_config,
+            remote_variables_dir=self.remote_variables_dir,
             retry_config=self.retry_config,
             run_pipeline_in_one_process=self.run_pipeline_in_one_process,
+            settings=self.settings.to_dict() if self.settings else self.settings,
             tags=self.tags,
             type=self.type.value if type(self.type) is not str else self.type,
-            updated_at=self.updated_at,
+            updated_at=updated_at,
             uuid=self.uuid,
+            variables_dir=self.variables_dir,
         )
 
         if self.variables is not None:
@@ -852,6 +874,10 @@ class Pipeline:
                 setattr(self, key, data.get(key))
                 should_save = True
 
+        if 'settings' in data:
+            self.settings = PipelineSettings.load(**(data.get('settings') or {}))
+            should_save = True
+
         blocks = data.get('blocks', [])
 
         if blocks:
@@ -888,6 +914,16 @@ class Pipeline:
                             self.extensions.get(extension_uuid, {}).get('blocks_by_uuid', {}),
                         ))
 
+            global_hooks = None
+            if len(arr) >= 1:
+                project = Project(self.repo_config)
+                if project.is_feature_enabled(FeatureUUID.GLOBAL_HOOKS):
+                    from mage_ai.data_preparation.models.global_hooks.models import (
+                        GlobalHooks,
+                    )
+
+                    global_hooks = GlobalHooks.load_from_file()
+
             for tup in arr:
                 key, blocks_arr, mapping = tup
                 widget = key == 'widgets'
@@ -901,6 +937,32 @@ class Pipeline:
                     block = mapping.get(block_data['uuid'])
                     if block is None:
                         continue
+
+                    if global_hooks:
+                        from mage_ai.data_preparation.models.global_hooks.models import (
+                            HookOperation,
+                            HookStage,
+                        )
+
+                        hooks = global_hooks.get_and_run_hooks(
+                            operation_resource=block,
+                            operation_types=[HookOperation.UPDATE_ANYWHERE],
+                            resource_type=EntityName.Block,
+                            stage=HookStage.BEFORE,
+                            payload=block_data,
+                        )
+
+                        if hooks:
+                            for hook in (hooks or []):
+                                output = hook.output
+                                if not output:
+                                    continue
+
+                                if output.get('payload'):
+                                    value = output.get('payload') or {}
+                                    if isinstance(value, dict):
+                                        block_data = merge_dict(block_data, value)
+
                     if 'content' in block_data:
                         from mage_ai.cache.block_action_object import (
                             BlockActionObjectCache,
@@ -1677,6 +1739,22 @@ class Pipeline:
             file_version_only=True,
         )
 
+    def should_save_trigger_in_code_automatically(self) -> bool:
+        from mage_ai.data_preparation.models.project import Project
+
+        if self.settings and \
+                self.settings.triggers and \
+                self.settings.triggers.save_in_code_automatically is not None:
+
+            return self.settings.triggers.save_in_code_automatically
+
+        project = Project(self.repo_config)
+
+        return project.pipelines and \
+            project.pipelines.settings and \
+            project.pipelines.settings.triggers and \
+            project.pipelines.settings.triggers.save_in_code_automatically
+
     async def save_async(
         self,
         block_type: str = None,

mage_ai/data_preparation/models/pipelines/integration_pipeline.py

@@ -272,7 +272,6 @@ class IntegrationPipeline(Pipeline):
                             sample_data_json = data.get('sample_data')
                             sample_data = pd.DataFrame.from_dict(json.loads(sample_data_json))
                             stream_id = data.get('stream_id')
-
                             variables = {
                                 f'output_sample_data_{stream_id}': sample_data,
                             }

mage_ai/data_preparation/models/pipelines/models.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class PipelineSettingsTriggers(BaseDataClass):
+    save_in_code_automatically: bool = None
+
+
+@dataclass
+class PipelineSettings(BaseDataClass):
+    triggers: PipelineSettingsTriggers = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('triggers', PipelineSettingsTriggers)

mage_ai/data_preparation/models/project/__init__.py

@@ -35,6 +35,8 @@ class Project():
         features = self.repo_config.features
 
         for uuid in FeatureUUID:
+            if FeatureUUID.GLOBAL_HOOKS == uuid:
+                continue
             key = uuid.value
             data[key] = features.get(key) if features else None
 
@@ -52,6 +54,10 @@ class Project():
     def remote_variables_dir(self) -> Dict:
         return self.repo_config.remote_variables_dir
 
+    @property
+    def pipelines(self) -> Dict:
+        return self.repo_config.pipelines
+
     def is_feature_enabled(self, feature_name: FeatureUUID) -> str:
         feature_enabled = self.repo_config.features.get(feature_name.value, False)
 

mage_ai/data_preparation/models/project/constants.py

@@ -5,6 +5,7 @@ class FeatureUUID(str, Enum):
     ADD_NEW_BLOCK_V2 = 'add_new_block_v2'
     COMPUTE_MANAGEMENT = 'compute_management'
     DATA_INTEGRATION_IN_BATCH_PIPELINE = 'data_integration_in_batch_pipeline'
+    GLOBAL_HOOKS = 'global_hooks'
     INTERACTIONS = 'interactions'
     LOCAL_TIMEZONE = 'display_local_timezone'
     NOTEBOOK_BLOCK_OUTPUT_SPLIT_VIEW = 'notebook_block_output_split_view'

mage_ai/data_preparation/models/project/models.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+from mage_ai.data_preparation.models.pipelines.models import PipelineSettings
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class ProjectPipelines(BaseDataClass):
+    settings: PipelineSettings = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('settings', PipelineSettings)