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

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,
@@ -626,9 +645,10 @@ class Pipeline:
             notification_config=self.notification_config,
             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,
         )
 
@@ -852,6 +872,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 +912,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 +935,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 +1737,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/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)

mage_ai/data_preparation/repo_manager.py

@@ -57,6 +57,7 @@ class RepoConfig:
         self.queue_config = None
         self.help_improve_mage = None
         self.openai_api_key = None
+        self._pipelines = None
         self.retry_config = None
         self.ldap_config = None
         self.s3_bucket = None
@@ -117,6 +118,7 @@ class RepoConfig:
             self.project_uuid = repo_config.get('project_uuid')
             self.help_improve_mage = repo_config.get('help_improve_mage')
             self.openai_api_key = repo_config.get('openai_api_key')
+            self.pipelines = repo_config.get('pipelines')
             self.retry_config = repo_config.get('retry_config')
 
             self.ldap_config = repo_config.get('ldap_config')
@@ -149,6 +151,22 @@ class RepoConfig:
         else:
             return get_metadata_path()
 
+    @property
+    def pipelines(self):
+        if isinstance(self._pipelines, dict):
+            self.pipelines = self._pipelines
+
+        return self._pipelines
+
+    @pipelines.setter
+    def pipelines(self, pipelines: Dict = None) -> None:
+        from mage_ai.data_preparation.models.project.models import ProjectPipelines
+
+        if isinstance(pipelines, dict):
+            self._pipelines = ProjectPipelines.load(**(pipelines or {}))
+        else:
+            self._pipelines = pipelines
+
     def to_dict(self, remote: bool = False) -> Dict:
         return dict(
             ai_config=self.ai_config,
@@ -160,6 +178,7 @@ class RepoConfig:
             help_improve_mage=self.help_improve_mage,
             notification_config=self.notification_config,
             openai_api_key=self.openai_api_key,
+            pipelines=self.pipelines.to_dict() if self.pipelines else self.pipelines,
             project_type=self.project_type,
             project_uuid=self.project_uuid,
             queue_config=self.queue_config,
@@ -179,7 +198,10 @@ class RepoConfig:
 
         for key, value in kwargs.items():
             data[key] = value
-            if hasattr(self, key):
+
+            if 'pipelines' == key:
+                self.pipelines = value
+            elif hasattr(self, key):
                 setattr(self, key, value)
 
         with open(self.metadata_path, 'w') as f:

mage_ai/data_preparation/templates/callbacks/base.jinja

@@ -12,5 +12,9 @@ def success_callback(parent_block_data, **kwargs):
 
 @callback('failure')
 def failure_callback(parent_block_data, **kwargs):
+    """
+    The error that caused the block to fail is passed in as keyword
+    argument __error.
+    """
     pass
 {% endblock %}

mage_ai/data_preparation/templates/data_exporters/streaming/kafka.yaml

@@ -1,6 +1,7 @@
 connector_type: kafka
 bootstrap_server: "localhost:9092"
 topic: topic_name
+api_version: 0.10.2
 
 # Uncomment the config below to use SSL config
 # security_protocol: "SSL"

mage_ai/data_preparation/templates/data_loaders/streaming/activemq.yaml

@@ -0,0 +1,6 @@
+connector_type: 'activemq'
+connection_host: 'localhost'
+connection_port: 61613
+queue_name: 'default'
+username: 'admin'
+password: 'admin'

mage_ai/data_preparation/templates/data_loaders/streaming/kafka.yaml

@@ -3,6 +3,7 @@ bootstrap_server: "localhost:9092"
 topic: topic_name
 consumer_group: unique_consumer_group
 include_metadata: false
+api_version: 0.10.2
 
 # Uncomment the config below to use SSL config
 # security_protocol: "SSL"

mage_ai/data_preparation/templates/data_loaders/streaming/nats.yaml

@@ -0,0 +1,20 @@
+connector_type: nats
+server_url: "nats://localhost:4222"
+subject: subject_name
+stream_name: stream_name
+
+# Optional: Batch size and timeout for reading messages
+batch_size: 100
+timeout: 5
+
+# Optional: TLS configuration
+use_tls: false
+# cafile: "path/to/ca.pem"
+# certfile: "path/to/client-cert.pem"
+# keyfile: "path/to/client-key.pem"
+
+# Optional: NKEY seed string
+# nkeys_seed_str: "SUAPAEYJQOQHJ4"
+
+# Optional: Consumer name for pull-based subscription
+consumer_name: "my_consumer"

mage_ai/data_preparation/templates/repo/io_config.yaml

@@ -68,8 +68,8 @@ default:
   # Pinot
   PINOT_HOST: hostname
   PINOT_PASSWORD: password
-  PINOT_PATH: /query
-  PINOT_PORT: 9000
+  PINOT_PATH: /query/sql
+  PINOT_PORT: 8000
   PINOT_SCHEME: http
   PINOT_USER: user
   # PostgresSQL

mage_ai/io/duckdb.py

@@ -83,7 +83,6 @@ class DuckDB(BaseSQL):
         cursor.execute(sql)
 
     def get_type(self, column: Series, dtype: str) -> str:
-        print(f"Type: {dtype}")
         if dtype in (
             PandasTypes.MIXED,
             PandasTypes.UNKNOWN_ARRAY,

mage_ai/orchestration/concurrency.py

@@ -3,7 +3,14 @@ from dataclasses import dataclass
 from mage_ai.shared.config import BaseConfig
 
 
+class OnLimitReached:
+    WAIT = 'wait'
+    SKIP = 'skip'
+
+
 @dataclass
 class ConcurrencyConfig(BaseConfig):
-    pipeline_run_limit: int = None
     block_run_limit: int = None
+    pipeline_run_limit: int = None
+    pipeline_run_limit_all_triggers: int = None
+    on_pipeline_run_limit_reached: OnLimitReached = OnLimitReached.WAIT

mage_ai/orchestration/db/models/schedules.py

@@ -3,6 +3,7 @@ import enum
 import traceback
 import uuid
 from datetime import datetime, timedelta, timezone
+from itertools import groupby
 from math import ceil
 from statistics import stdev
 from typing import Dict, List
@@ -656,7 +657,7 @@ class PipelineRun(BaseModel):
     def pipeline_type(self) -> PipelineType:
         pipeline = Pipeline.get(self.pipeline_uuid, check_if_exists=True)
 
-        return self.pipeline.type if pipeline is not None else None
+        return pipeline.type if pipeline is not None else None
 
     @property
     def logs(self):
@@ -689,7 +690,7 @@ class PipelineRun(BaseModel):
     def pipeline_tags(self):
         pipeline = Pipeline.get(self.pipeline_uuid, check_if_exists=True)
 
-        return self.pipeline.tags if pipeline is not None else []
+        return pipeline.tags if pipeline is not None else []
 
     def executable_block_runs(
         self,
@@ -858,6 +859,26 @@ class PipelineRun(BaseModel):
             query = query.options(joinedload(PipelineRun.block_runs))
         return query.all()
 
+    @classmethod
+    @safe_db_query
+    def active_runs_for_pipelines_grouped(
+        self,
+        pipeline_uuids: List[str],
+        include_block_runs: bool = False,
+    ) -> Dict[str, List['PipelineRun']]:
+        """
+        Get a dictionary of active pipeline runs grouped by pipeline uuid.
+        """
+
+        active_runs = self.active_runs_for_pipelines(
+            pipeline_uuids,
+            include_block_runs=include_block_runs,
+        )
+        grouped = {}
+        for key, runs in groupby(active_runs, key=lambda x: x.pipeline_uuid):
+            grouped[key] = list(runs)
+        return grouped
+
     @classmethod
     @safe_db_query
     def batch_update_status(self, pipeline_run_ids: List[int], status):