mage-ai 0.9.45__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/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/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/orchestration/db/models/schedules.py

@@ -657,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):
@@ -690,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,

mage_ai/orchestration/queue/process_queue.py

@@ -116,7 +116,15 @@ class ProcessQueue(Queue):
             if job_client_id != self.client_id and self.redis_client.get(job_client_id):
                 return True
         job = self.job_dict.get(job_id)
-        return job is not None and (job == JobStatus.QUEUED or isinstance(job, int))
+        return (
+            job is not None and
+            (
+                # In queue
+                (job == JobStatus.QUEUED and not self.queue.empty()) or
+                # Running
+                isinstance(job, int)
+            )
+        )
 
     def kill_job(self, job_id: str):
         """

mage_ai/orchestration/triggers/api.py

@@ -22,18 +22,21 @@ def trigger_pipeline(
     error_on_failure: bool = False,
     poll_interval: float = DEFAULT_POLL_INTERVAL,
     poll_timeout: Optional[float] = None,
+    schedule_name: str = None,
     verbose: bool = True,
+    _should_schedule: bool = False,  # For internal use only (e.g. running hooks from notebook).
 ) -> PipelineRun:
     if variables is None:
         variables = {}
     pipeline = Pipeline.get(pipeline_uuid)
 
-    pipeline_schedule = __fetch_or_create_pipeline_schedule(pipeline)
+    pipeline_schedule = __fetch_or_create_pipeline_schedule(pipeline, schedule_name=schedule_name)
 
     pipeline_run = create_and_start_pipeline_run(
         pipeline,
         pipeline_schedule,
         dict(variables=variables),
+        should_schedule=_should_schedule,
     )
 
     if check_status:
@@ -48,9 +51,14 @@ def trigger_pipeline(
     return pipeline_run
 
 
-def __fetch_or_create_pipeline_schedule(pipeline: Pipeline) -> PipelineSchedule:
+def __fetch_or_create_pipeline_schedule(
+    pipeline: Pipeline,
+    schedule_name: str = None,
+) -> PipelineSchedule:
+    if schedule_name is None:
+        schedule_name = TRIGGER_NAME_FOR_TRIGGER_CREATED_FROM_CODE
+
     pipeline_uuid = pipeline.uuid
-    schedule_name = TRIGGER_NAME_FOR_TRIGGER_CREATED_FROM_CODE
     schedule_type = ScheduleType.API
 
     pipeline_schedule = PipelineSchedule.repo_query.filter(

mage_ai/orchestration/triggers/constants.py

@@ -1,3 +1,4 @@
 DEFAULT_POLL_INTERVAL = 60
 TRIGGER_NAME_FOR_GLOBAL_DATA_PRODUCT = 'GLOBAL_DATA_PRODUCT_TRIGGER'
+TRIGGER_NAME_FOR_GLOBAL_HOOK = 'GLOBAL_HOOK_TRIGGER'
 TRIGGER_NAME_FOR_TRIGGER_CREATED_FROM_CODE = 'API_TRIGGER_CREATED_FROM_CODE'

mage_ai/server/active_kernel.py

@@ -3,8 +3,13 @@ from typing import Dict
 from jupyter_client import KernelClient, KernelManager
 from jupyter_client.kernelspec import NoSuchKernel
 
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.server.kernels import DEFAULT_KERNEL_NAME, KernelName, kernel_managers
 from mage_ai.server.logger import Logger
+from mage_ai.services.spark.constants import ComputeServiceUUID
+from mage_ai.services.spark.utils import get_compute_service
+from mage_ai.services.ssh.aws.emr.utils import cluster_info_from_tunnel
 
 logger = Logger().new_server_logger(__name__)
 
@@ -68,10 +73,38 @@ def switch_active_kernel(
             from mage_ai.cluster_manager.aws.emr_cluster_manager import (
                 emr_cluster_manager,
             )
-            emr_cluster_manager.set_active_cluster(
-                auto_selection=True,
-                emr_config=emr_config,
-            )
+
+            should_set_active = True
+            auto_creation = True
+            cluster_id = None
+            project = Project()
+
+            if project.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
+                if ComputeServiceUUID.AWS_EMR == get_compute_service(
+                    emr_config=emr_config,
+                    kernel_name=kernel_name,
+                ):
+                    auto_creation = False
+                    should_set_active = False
+
+                    cluster_info = cluster_info_from_tunnel()
+
+                    if cluster_info:
+                        from mage_ai.services.compute.models import ComputeService
+
+                        cluster_id = cluster_info.get('id') or None
+                        compute_service = ComputeService(project=project)
+                        cluster = compute_service.get_cluster_details(cluster_id=cluster_id)
+
+                        should_set_active = cluster.has_dns_name if cluster else False
+
+            if should_set_active:
+                emr_cluster_manager.set_active_cluster(
+                    auto_creation=auto_creation,
+                    auto_selection=True,
+                    cluster_id=cluster_id,
+                    emr_config=emr_config,
+                )
     except NoSuchKernel as e:
         if kernel_name == KernelName.PYSPARK:
             raise Exception(

mage_ai/server/api/downloads.py

@@ -1,3 +1,8 @@
+import os
+import tempfile
+import zipfile
+
+import jwt
 from tornado import gen, iostream
 
 from mage_ai.api.utils import authenticate_client_and_token
@@ -6,7 +11,8 @@ from mage_ai.data_preparation.models.variable import VariableType
 from mage_ai.orchestration.db.models.oauth import Oauth2Application
 from mage_ai.orchestration.db.models.schedules import PipelineRun
 from mage_ai.server.api.base import BaseHandler
-from mage_ai.settings import REQUIRE_USER_AUTHENTICATION
+from mage_ai.settings import JWT_DOWNLOAD_SECRET, REQUIRE_USER_AUTHENTICATION
+from mage_ai.settings.repo import get_repo_path
 
 
 class ApiDownloadHandler(BaseHandler):
@@ -62,3 +68,72 @@ class ApiDownloadHandler(BaseHandler):
                     await self.flush()
                     # Sleep for a nanosecond so other handlers can run and avoid blocking
                     await gen.sleep(0.000000001)
+
+
+class ApiResourceDownloadHandler(BaseHandler):
+
+    def get(self, token):
+        try:
+            decoded_payload = jwt.decode(token, JWT_DOWNLOAD_SECRET, algorithms=['HS256'])
+
+            file_name = decoded_payload['file_name']
+            file_list = decoded_payload['file_list']
+            self.ignore_folder_structure = decoded_payload['ignore_folder_structure']
+
+            self.abs_repo_path = os.path.abspath(get_repo_path())
+
+            relative_file_list = list(map(self.relative_path_mapping, file_list))
+
+            try:
+                file_pointer = self.get_file_pointer(file_list, relative_file_list)
+
+                while True:
+                    _buffer = file_pointer.read(4096)
+                    if not _buffer:
+                        break
+                    self.write(_buffer)
+            except Exception as e:
+                self.set_status(400)
+                self.write(f'Error fetching file {file_name}.\n{e}')
+            finally:
+                file_pointer.close()
+
+            self.set_header('Content-Type', 'application/force-download')
+            self.set_header('Content-Disposition', f'attachment; filename={file_name}')
+            self.flush()
+        except jwt.exceptions.ExpiredSignatureError:
+            self.set_status(400)
+            self.write('Download token is expired.')
+        except (jwt.exceptions.InvalidSignatureError, jwt.exceptions.DecodeError):
+            self.set_status(400)
+            self.write('Download token is invalid.')
+        except ValueError as e:
+            self.set_status(400)
+            self.write(f'Attepmt at fetching file outside of project folder: {e}')
+
+    # file pointer points to either a singular file or a temporary zip
+    def get_file_pointer(self, file_list, relative_file_list):
+        if len(file_list) == 1:
+            return open(file_list[0])
+        return self.zip_files(file_list, relative_file_list)
+
+    # creates a temporary zip and returns the (open) file pointer
+    def zip_files(self, file_list, relative_file_list):
+        zip_file = tempfile.NamedTemporaryFile(suffix='.zip', mode='w+b')
+        with zipfile.ZipFile(zip_file, 'w') as zipf:
+            for path, relative in zip(file_list, relative_file_list):
+                zipf.write(path, relative)
+        zip_file.seek(0)  # set cursor to start of file to prepare for content extraction
+        return zip_file
+
+    def relative_path_mapping(self, path):
+        abs_path = os.path.abspath(path)
+        common_ground = os.path.commonpath([self.abs_repo_path, abs_path])
+
+        # trying to access files outside of the project folder
+        if common_ground != self.abs_repo_path:
+            raise ValueError(abs_path)
+
+        return (os.path.basename(abs_path)
+                if self.ignore_folder_structure
+                else os.path.relpath(abs_path, common_ground))

mage_ai/server/constants.py

@@ -12,4 +12,4 @@ DATAFRAME_OUTPUT_SAMPLE_COUNT = 10
 # Dockerfile depends on it because it runs ./scripts/install_mage.sh and uses
 # the last line to determine the version to install.
 VERSION = \
-'0.9.45'
+'0.9.46'