mage-ai 0.9.46__py3-none-any.whl → 0.9.48__py3-none-any.whl

mage_ai/data_preparation/models/global_hooks/models.py

@@ -1,22 +1,23 @@
 import hashlib
 import os
-from collections.abc import Iterable
 from dataclasses import dataclass, field, make_dataclass
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 import yaml
 
 from mage_ai.api.operations.constants import OperationType
 from mage_ai.api.resources.BaseResource import BaseResource
 from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.block.utils import fetch_input_variables
 from mage_ai.data_preparation.models.global_hooks.constants import (
     RESOURCE_TYPES,
     HookOutputKey,
 )
+from mage_ai.data_preparation.models.global_hooks.predicates import HookPredicate
 from mage_ai.data_preparation.models.global_hooks.utils import extract_valid_data
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db.models.schedules import PipelineRun
@@ -24,6 +25,7 @@ from mage_ai.orchestration.triggers.api import trigger_pipeline
 from mage_ai.orchestration.triggers.constants import TRIGGER_NAME_FOR_GLOBAL_HOOK
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.array import find, find_index, flatten
+from mage_ai.shared.environments import is_debug, is_test
 from mage_ai.shared.hash import (
     dig,
     extract,
@@ -75,6 +77,7 @@ class HookStrategies(BaseDataClass):
 
 @dataclass
 class HookRunSettings(BaseDataClass):
+    asynchronous: bool = False
     with_trigger: bool = False
 
 
@@ -109,11 +112,6 @@ class HookOutputSettings(BaseDataClass):
         return merge_dict(ignore_keys_with_blank_values(super().to_dict(**kwargs)), data)
 
 
-@dataclass
-class HookPredicate(BaseDataClass):
-    resource: Dict = field(default_factory=dict)
-
-
 @dataclass
 class MetadataUser(BaseDataClass):
     id: int = None
@@ -144,7 +142,7 @@ class Hook(BaseDataClass):
     output: Dict = field(default=None)
     output_settings: List[HookOutputSettings] = None
     pipeline_settings: Dict = field(default=None)
-    predicates: List[List[HookPredicate]] = None
+    predicate: HookPredicate = None
     resource_type: EntityName = None
     run_settings: HookRunSettings = None
     stages: List[HookStage] = None
@@ -156,6 +154,7 @@ class Hook(BaseDataClass):
         self._pipeline = None
 
         self.serialize_attribute_class('metadata', HookMetadata)
+        self.serialize_attribute_class('predicate', HookPredicate)
         self.serialize_attribute_class('run_settings', HookRunSettings)
         self.serialize_attribute_class('status', HookStatus)
         self.serialize_attribute_class('strategies', HookStrategies)
@@ -165,16 +164,6 @@ class Hook(BaseDataClass):
         self.serialize_attribute_enums('conditions', HookCondition)
         self.serialize_attribute_enums('stages', HookStage)
 
-        if self.predicates and isinstance(self.predicates, list):
-            rows = []
-            for predicates in self.predicates:
-                row = []
-                if predicates and isinstance(predicates, list):
-                    for predicate in predicates:
-                        row.append(HookPredicate.load(**predicate))
-                rows.append(row)
-            self.predicates = rows
-
     def to_dict(
         self,
         include_all: bool = False,
@@ -185,6 +174,7 @@ class Hook(BaseDataClass):
         arr = [
             'conditions',
             'metadata',
+            'predicate',
             'run_settings',
             'stages',
             'strategies',
@@ -211,20 +201,6 @@ class Hook(BaseDataClass):
         if self.pipeline_settings:
             data['pipeline'] = self.pipeline_settings
 
-        if self.predicates:
-            rows = []
-            for predicates in self.predicates:
-                if predicates:
-                    row = []
-                    for predicate in predicates:
-                        row.append(predicate.to_dict())
-
-                    if len(row) >= 1:
-                        rows.append(row)
-
-            if len(rows) >= 1:
-                data['predicates'] = rows
-
         if include_snapshot_validation:
             key = 'snapshot_valid'
             if data.get('metadata'):
@@ -360,7 +336,10 @@ class Hook(BaseDataClass):
         check_status: bool = True,
         error_on_failure: bool = True,
         poll_timeout: int = None,
+        resource_id: Union[int, str] = None,
+        resource_parent_id: Union[int, str] = None,
         should_schedule: bool = False,
+        user: Dict = None,
         with_trigger: bool = False,
         **kwargs,
     ) -> None:
@@ -375,7 +354,10 @@ class Hook(BaseDataClass):
                 payload=kwargs.get('payload'),
                 query=kwargs.get('query'),
                 resource=kwargs.get('resource'),
+                resource_id=resource_id,
+                resource_parent_id=resource_parent_id,
                 resources=kwargs.get('resources'),
+                user=user,
             ))
             variables = merge_dict(
                 self.pipeline_settings.get('variables') or {},
@@ -388,11 +370,14 @@ class Hook(BaseDataClass):
             )
 
             pipeline_run = None
+
+            asynchronous = self.run_settings.asynchronous if self.run_settings else False
+
             if with_trigger or (self.run_settings and self.run_settings.with_trigger):
                 pipeline_run = trigger_pipeline(
                     self.pipeline.uuid,
                     variables=variables,
-                    check_status=check_status,
+                    check_status=check_status and not asynchronous,
                     error_on_failure=error_on_failure,
                     poll_interval=1,
                     poll_timeout=poll_timeout,
@@ -400,6 +385,10 @@ class Hook(BaseDataClass):
                     verbose=True,
                     _should_schedule=should_schedule,
                 )
+            elif asynchronous:
+                # TODO: invoking the below method will still block the current
+                # operation from completing
+                PipelineExecutor(self.pipeline).execute(global_vars=variables, update_status=False)
             else:
                 self.pipeline.execute_sync(global_vars=variables, update_status=False)
 
@@ -417,13 +406,26 @@ class Hook(BaseDataClass):
                 elif HookStrategy.CONTINUE in self.strategies:
                     self.status.strategy = HookStrategy.CONTINUE
 
+            if is_debug() or is_test():
+                print(f'[ERROR] Hook.run: {err}')
+
     def should_run(
         self,
         operation_types: List[HookOperation],
         resource_type: EntityName,
         stage: HookStage,
         conditions: List[HookCondition] = None,
+        error: Dict = None,
+        meta: Dict = None,
+        metadata: Dict = None,
         operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline] = 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:
         if self.operation_type not in operation_types:
             return False
@@ -440,7 +442,19 @@ class Hook(BaseDataClass):
         if not self.__validate_snapshot():
             return False
 
-        if not self.__matches_any_predicate(operation_resource):
+        if not self.__matches_predicate(
+            operation_resource,
+            error=error,
+            meta=meta,
+            metadata=metadata,
+            payload=payload,
+            query=query,
+            resource=resource,
+            resource_id=resource_id,
+            resource_parent_id=resource_parent_id,
+            resources=resources,
+            user=user,
+        ):
             return False
 
         return True
@@ -490,52 +504,37 @@ class Hook(BaseDataClass):
 
         return any([condition in (self.conditions or []) for condition in conditions])
 
-    def __matches_any_predicate(
+    def __matches_predicate(
         self,
         operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
+        error: 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:
-        if not operation_resource or not self.predicates:
+        if not operation_resource or not self.predicate:
             return True
 
-        return any([all([self.__validate_predicate(
-            predicate,
+        return self.predicate.validate(
             operation_resource,
-        ) for predicate in predicates]) for predicates in self.predicates])
-
-    def __validate_predicate(
-        self,
-        predicate: HookPredicate,
-        operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline],
-    ) -> bool:
-        if not predicate.resource or len(predicate.resource) == 0:
-            return True
-
-        def _validate_resource(
-            resource: Union[BaseResource, Block, Dict, Pipeline],
-            predicate=predicate,
-        ) -> bool:
-            model = resource
-            if isinstance(resource, BaseResource):
-                model = resource.model
-
-            def _equals(
-                key: str,
-                value: Any,
-                model=model,
-            ) -> bool:
-                if isinstance(model, dict):
-                    return model.get(key) == value
-
-                return hasattr(model, key) and getattr(model, key) == value
-
-            check = all([_equals(key, value) for key, value in predicate.resource.items()])
-
-            return check
-
-        if isinstance(operation_resource, Iterable) and not isinstance(operation_resource, dict):
-            return all([_validate_resource(res) for res in operation_resource])
-
-        return _validate_resource(operation_resource)
+            error=error,
+            hook=self.to_dict(include_all=True),
+            meta=meta,
+            metadata=metadata,
+            payload=payload,
+            query=query,
+            resource=resource,
+            resource_id=resource_id,
+            resource_parent_id=resource_parent_id,
+            resources=resources,
+            user=user,
+        )
 
 
 def __build_global_hook_resource_fields() -> List[Tuple]:
@@ -838,22 +837,52 @@ class GlobalHooks(BaseDataClass):
         resource_type: EntityName,
         stage: HookStage,
         conditions: List[HookCondition] = None,
+        error: Dict = None,
+        meta: Dict = None,
+        metadata: Dict = None,
         operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline] = None,
+        payload: Dict = None,
+        query: Dict = None,
+        resource: Dict = None,
+        resource_id: Union[int, str] = None,
+        resource_parent_id: Union[int, str] = None,
+        resources: List[Dict] = None,
+        user: Dict = None,
     ) -> List[Hook]:
         def _filter(
-            hook: Hook,
+            hook,
             conditions=conditions,
+            error=error,
+            meta=meta,
+            metadata=metadata,
             operation_resource=operation_resource,
             operation_types=operation_types,
+            payload=payload,
+            query=query,
+            resource=resource,
+            resource_id=resource_id,
+            resource_parent_id=resource_parent_id,
             resource_type=resource_type,
+            resources=resources,
             stage=stage,
+            user=user,
         ) -> bool:
             return hook.should_run(
+                operation_types,
+                resource_type,
+                stage,
                 conditions=conditions,
+                error=error,
+                meta=meta,
+                metadata=metadata,
                 operation_resource=operation_resource,
-                operation_types=operation_types,
-                resource_type=resource_type,
-                stage=stage,
+                payload=payload,
+                query=query,
+                resource=resource,
+                resource_id=resource_id,
+                resource_parent_id=resource_parent_id,
+                resources=resources,
+                user=user,
             )
 
         return list(filter(
@@ -861,7 +890,11 @@ class GlobalHooks(BaseDataClass):
             self.hooks(),
         ))
 
-    def run_hooks(self, hooks: List[Hook], **kwargs) -> List[Hook]:
+    def run_hooks(
+        self,
+        hooks: List[Hook],
+        **kwargs,
+    ) -> List[Hook]:
         hooks_by_pipeline = {}
         for hook in hooks:
             if not hook.pipeline_settings or not hook.pipeline_settings.get('uuid'):
@@ -889,21 +922,52 @@ class GlobalHooks(BaseDataClass):
         resource_type: EntityName,
         stage: HookStage,
         conditions: List[HookCondition] = None,
+        error: Dict = None,
+        meta: Dict = None,
+        metadata: Dict = None,
         operation_resource: Union[BaseResource, Block, Dict, List[BaseResource], Pipeline] = None,
-        **kwargs,
+        payload: Dict = None,
+        query: Dict = None,
+        resource: Dict = None,
+        resource_id: Union[int, str] = None,
+        resource_parent_id: Union[int, str] = None,
+        resources: List[Dict] = None,
+        user: Dict = None,
     ) -> List[Hook]:
         hooks = self.get_hooks(
-            operation_types=operation_types,
-            resource_type=resource_type,
-            stage=stage,
+            operation_types,
+            resource_type,
+            stage,
             conditions=conditions,
+            error=error,
+            meta=meta,
+            metadata=metadata,
             operation_resource=operation_resource,
+            payload=payload,
+            query=query,
+            resource=resource,
+            resource_id=resource_id,
+            resource_parent_id=resource_parent_id,
+            resources=resources,
+            user=user,
         )
 
         if not hooks:
             return None
 
-        return self.run_hooks(hooks, **kwargs)
+        return self.run_hooks(
+            hooks,
+            error=error,
+            meta=meta,
+            metadata=metadata,
+            payload=payload,
+            query=query,
+            resource=resource,
+            resource_id=resource_id,
+            resource_parent_id=resource_parent_id,
+            resources=resources,
+            user=user,
+        )
 
     def save(self, file_path: str = None) -> None:
         if not file_path: