mage-ai 0.9.69__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/data_preparation/models/pipelines/integration_pipeline.py

@@ -262,8 +262,7 @@ class IntegrationPipeline(Pipeline):
                 output = proc.stdout.decode()
                 print_logs_from_output(output, config=config)
 
-                pipeline = Pipeline(self.uuid)
-                block = pipeline.get_block(self.data_loader.uuid)
+                block = self.get_block(self.data_loader.uuid)
 
                 for line in output.split('\n'):
                     try:

mage_ai/data_preparation/models/pipelines/seed.py

@@ -30,7 +30,7 @@ def seed(
     pipeline_models = [Pipeline.get(
         uuid,
         repo_path=repo_path,
-    ) for uuid in Pipeline.get_all_pipelines(repo_path=repo_path)]
+    ) for uuid in Pipeline.get_all_pipelines(repo_path)]
     for i1 in range(pipelines - len(pipeline_models)):
         i = i1 + len(pipeline_models)
         model = Pipeline.create(name=f'Pipeline {i}', repo_path=repo_path)

mage_ai/data_preparation/models/project/__init__.py

@@ -1,5 +1,5 @@
 import os
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import aiohttp
 
@@ -18,16 +18,36 @@ from mage_ai.shared.environments import is_debug
 from mage_ai.shared.hash import dig
 
 
-class Project():
-    def __init__(self, repo_config=None, repo_path: str = None, root_project: bool = False):
+class Project:
+    def __init__(
+        self,
+        repo_config=None,
+        repo_path: Optional[str] = None,
+        context_data: Dict = None,
+        root_project: bool = False,
+        user=None,
+    ):
+        # print(f'Project context_data {id(context_data)} {context_data}')
+        if context_data is None:
+            context_data = dict()
+        self.context_data = context_data
         self.root_project = root_project
-        self.repo_path = repo_path or get_repo_path(root_project=self.root_project)
+        self.repo_path = repo_path or get_repo_path(
+            context_data=context_data,
+            root_project=self.root_project,
+            user=user,
+        )
+        self.user = user
 
         self.name = os.path.basename(self.repo_path)
         self.settings = None
 
         if not root_project and project_platform_activated():
-            self.settings = active_project_settings(get_default=True)
+            self.settings = active_project_settings(
+                context_data=context_data,
+                get_default=True,
+                user=user,
+            )
             if self.settings and self.settings.get('uuid'):
                 self.name = self.settings.get('uuid')
 
@@ -55,12 +75,13 @@ class Project():
     @property
     def workspace_config_defaults(self) -> Dict:
         config = self.repo_config.workspace_config_defaults or {}
-        cluster_type = get_cluster_type()
+        cluster_type = get_cluster_type(repo_path=self.repo_path)
         try:
             if cluster_type == ClusterType.K8S:
                 from mage_ai.cluster_manager.kubernetes.workload_manager import (
                     WorkloadManager,
                 )
+
                 workload_manager = WorkloadManager(os.getenv(KUBE_NAMESPACE))
                 k8s_default_values = workload_manager.get_default_values()
 
@@ -108,7 +129,9 @@ class Project():
 
         if project_platform_activated() and self.__repo_config_root_project:
             settings = platform_settings()
-            if settings.get('features') and (settings.get('features') or {}).get('override'):
+            if settings.get('features') and (settings.get('features') or {}).get(
+                'override'
+            ):
                 features.update(self.features_override)
 
         for uuid in FeatureUUID:
@@ -122,7 +145,9 @@ class Project():
     def features_override(self) -> Dict:
         if project_platform_activated() and self.__repo_config_root_project:
             settings = self.platform_settings()
-            if settings.get('features') and (settings.get('features') or {}).get('override'):
+            if settings.get('features') and (settings.get('features') or {}).get(
+                'override'
+            ):
                 return self.__repo_config_root_project.features or {}
 
         return {}
@@ -144,12 +169,25 @@ class Project():
         return self.repo_config.pipelines
 
     @classmethod
-    def is_feature_enabled_in_root_or_active_project(self, feature_name: FeatureUUID) -> bool:
-        if self(root_project=True).is_feature_enabled(feature_name):
+    def is_feature_enabled_in_root_or_active_project(
+        cls,
+        feature_name: FeatureUUID,
+        context_data: Dict = None,
+        user=None,
+    ) -> bool:
+        if cls(
+            context_data=context_data,
+            root_project=True,
+            user=user,
+        ).is_feature_enabled(feature_name):
             return True
 
         if project_platform_activated():
-            return self(root_project=False).is_feature_enabled(feature_name)
+            return cls(
+                context_data=context_data,
+                root_project=False,
+                user=user,
+            ).is_feature_enabled(feature_name)
 
         return False
 
@@ -161,17 +199,27 @@ class Project():
         if self.settings:
             query_arr = dig(self.settings, ['database', 'query', key])
             if query_arr:
-                return [os.path.join(*[part for part in [
-                    os.path.dirname(get_repo_path(root_project=True)),
-                    query_alias,
-                ] if len(part) >= 1]) for query_alias in query_arr] + [
-                    get_repo_path(root_project=False)
-                ]
+                return [
+                    os.path.join(
+                        *[
+                            part
+                            for part in [
+                                os.path.dirname(get_repo_path(root_project=True)),
+                                query_alias,
+                            ]
+                            if len(part) >= 1
+                        ]
+                    )
+                    for query_alias in query_arr
+                ] + [get_repo_path(root_project=False)]
 
         return [self.repo_path]
 
     def projects(self) -> Dict:
-        return project_platform_settings(mage_projects_only=True)
+        return project_platform_settings(
+            context_data=self.context_data,
+            mage_projects_only=True
+        )
 
     def is_feature_enabled(self, feature_name: FeatureUUID) -> bool:
         feature_enabled = self.features.get(feature_name.value, False)

mage_ai/data_preparation/models/project/constants.py

@@ -3,6 +3,7 @@ from enum import Enum
 
 class FeatureUUID(str, Enum):
     ADD_NEW_BLOCK_V2 = 'add_new_block_v2'
+    AUTOMATIC_KERNEL_CLEANUP = 'automatic_kernel_cleanup'
     CODE_BLOCK_V2 = 'code_block_v2'
     COMMAND_CENTER = 'command_center'
     COMPUTE_MANAGEMENT = 'compute_management'
@@ -14,3 +15,4 @@ class FeatureUUID(str, Enum):
     LOCAL_TIMEZONE = 'display_local_timezone'
     NOTEBOOK_BLOCK_OUTPUT_SPLIT_VIEW = 'notebook_block_output_split_view'
     OPERATION_HISTORY = 'operation_history'
+    POLARS = 'polars'

mage_ai/data_preparation/models/triggers/__init__.py

@@ -2,8 +2,8 @@ import enum
 import os
 import traceback
 from dataclasses import dataclass, field
-from datetime import datetime
-from typing import Dict, List
+from datetime import datetime, timezone
+from typing import Dict, List, Optional
 
 import yaml
 from croniter import croniter
@@ -14,6 +14,7 @@ from mage_ai.shared.config import BaseConfig
 from mage_ai.shared.constants import VALID_ENVS
 from mage_ai.shared.hash import index_by
 from mage_ai.shared.io import safe_write
+from mage_ai.shared.yaml import load_yaml
 
 TRIGGER_FILE_NAME = 'triggers.yaml'
 
@@ -30,9 +31,9 @@ class ScheduleType(str, enum.Enum):
 
 
 SCHEDULE_TYPE_TO_LABEL = {
-  ScheduleType.API: 'API',
-  ScheduleType.EVENT: 'Event',
-  ScheduleType.TIME: 'Schedule',
+    ScheduleType.API: 'API',
+    ScheduleType.EVENT: 'Event',
+    ScheduleType.TIME: 'Schedule',
 }
 
 
@@ -66,10 +67,12 @@ class Trigger(BaseConfig):
     status: ScheduleStatus = ScheduleStatus.INACTIVE
     last_enabled_at: datetime = None
     variables: Dict = field(default_factory=dict)
-    sla: int = None     # in seconds
+    sla: int = None  # in seconds
     settings: Dict = field(default_factory=dict)
     envs: List = field(default_factory=list)
     repo_path: str = None
+    description: Optional[str] = None
+    token: Optional[str] = None
 
     def __post_init__(self):
         if self.schedule_type and type(self.schedule_type) is str:
@@ -77,15 +80,19 @@ class Trigger(BaseConfig):
         if self.status and type(self.status) is str:
             self.status = ScheduleStatus(self.status)
         if any(env not in VALID_ENVS for env in self.envs):
-            raise Exception(f'Please provide valid env values inside {list(VALID_ENVS)}.')
+            raise Exception(
+                f'Please provide valid env values inside {list(VALID_ENVS)}.'
+            )
 
     @property
     def has_valid_schedule_interval(self) -> bool:
         # Check if trigger has valid cron expression
-        if self.schedule_interval is not None and \
-            self.schedule_type == ScheduleType.TIME and \
-            self.schedule_interval not in [e.value for e in ScheduleInterval] and \
-                not croniter.is_valid(self.schedule_interval):
+        if (
+            self.schedule_interval is not None
+            and self.schedule_type == ScheduleType.TIME
+            and self.schedule_interval not in [e.value for e in ScheduleInterval]
+            and not croniter.is_valid(self.schedule_interval)
+        ):
             return False
 
         return True
@@ -97,25 +104,43 @@ class Trigger(BaseConfig):
             name=self.name,
             pipeline_uuid=self.pipeline_uuid,
             schedule_interval=self.schedule_interval,
-            schedule_type=self.schedule_type.value if self.schedule_type else self.schedule_type,
+            schedule_type=self.schedule_type.value
+            if self.schedule_type
+            else self.schedule_type,
             settings=self.settings,
             sla=self.sla,
+            description=self.description,
+            token=self.token,
             start_time=self.start_time,
             status=self.status.value if self.status else self.status,
             variables=self.variables,
         )
 
 
-def get_triggers_file_path(pipeline_uuid: str) -> str:
-    pipeline_path = os.path.join(get_repo_path(), PIPELINES_FOLDER, pipeline_uuid)
+"""
+Cache the triggers config by pipeline uuid
+  Key: trigger_file_path
+  Value: {"updated_at": xxx, "triggers": []}
+"""
+triggers_cache = dict()
+
+
+def get_triggers_file_path(
+    pipeline_uuid: str,
+    repo_path: str = None
+) -> str:
+    pipeline_path = os.path.join(repo_path or get_repo_path(), PIPELINES_FOLDER, pipeline_uuid)
     trigger_file_path = os.path.join(pipeline_path, TRIGGER_FILE_NAME)
     return trigger_file_path
 
 
-def load_triggers_file_content(pipeline_uuid: str) -> str:
+def load_triggers_file_content(
+    pipeline_uuid: str,
+    repo_path: str = None,
+) -> str:
     content = None
 
-    trigger_file_path = get_triggers_file_path(pipeline_uuid)
+    trigger_file_path = get_triggers_file_path(pipeline_uuid, repo_path=repo_path)
     if os.path.exists(trigger_file_path):
         with open(trigger_file_path) as fp:
             content = fp.read()
@@ -133,19 +158,83 @@ def load_triggers_file_data(pipeline_uuid: str) -> Dict:
     return data
 
 
-def get_triggers_by_pipeline(pipeline_uuid: str) -> List[Trigger]:
-    trigger_file_path = get_triggers_file_path(pipeline_uuid)
+def get_triggers_by_pipeline_with_cache(
+    pipeline_uuid: str,
+    repo_path: str = None,
+):
+    """
+    Retrieve triggers for a given pipeline, utilizing a cache to avoid unnecessary file reads.
+
+    This function checks if the trigger file for the specified pipeline has been modified since
+    the last time it was cached. If the file has not been modified, the cached triggers are
+    returned. If the file has been modified or is not in the cache, it reads the file, updates
+    the cache, and returns the triggers.
+
+    Args:
+        pipeline_uuid (str): The unique identifier of the pipeline.
+        repo_path (str, optional): The path to the repository containing the pipeline. Defaults to
+            None.
+
+    Returns:
+        tuple: A tuple containing:
+            - triggers (list): A list of triggers for the specified pipeline.
+            - from_cache (bool): A boolean indicating whether the triggers were retrieved from the
+                cache.
+
+    Raises:
+        Exception: If there is an error loading the triggers file or parsing the triggers
+            configuration.
+    """
+    trigger_file_path = get_triggers_file_path(
+        pipeline_uuid,
+        repo_path=repo_path,
+    )
 
     if not os.path.exists(trigger_file_path):
-        return []
+        return [], False
+
+    trigger_file_updated_at = datetime.fromtimestamp(
+        os.path.getmtime(trigger_file_path),
+        tz=timezone.utc,
+    )
+
+    if trigger_file_path in triggers_cache:
+        # Try reading triggers from cache first
+        triggers_config = triggers_cache[trigger_file_path]
+        last_updated_at = triggers_config.get('updated_at')
+        triggers = triggers_config.get('triggers')
+        if last_updated_at and triggers and last_updated_at >= trigger_file_updated_at:
+            # Trigger file not modified since last time
+            return triggers, True
 
     try:
-        content = load_triggers_file_content(pipeline_uuid)
-        triggers = load_trigger_configs(content, pipeline_uuid=pipeline_uuid)
+        content = load_triggers_file_content(
+            pipeline_uuid,
+            repo_path=repo_path,
+        )
+        triggers = load_trigger_configs(
+            content,
+            pipeline_uuid=pipeline_uuid,
+            repo_path=repo_path,
+        )
+        # Update the cache with the new triggers and their update time
+        triggers_cache[trigger_file_path] = {
+            'updated_at': trigger_file_updated_at,
+            'triggers': triggers,
+        }
     except Exception:
         traceback.print_exc()
         triggers = []
 
+    return triggers, False
+
+
+def get_triggers_by_pipeline(
+    pipeline_uuid: str,
+    repo_path: str = None,
+) -> List[Trigger]:
+    # For backward compatibility
+    triggers, _ = get_triggers_by_pipeline_with_cache(pipeline_uuid, repo_path=repo_path)
     return triggers
 
 
@@ -183,7 +272,8 @@ def build_triggers(
         except Exception as e:
             if raise_exception:
                 raise Exception(
-                    f'Failed to parse trigger config {trigger_config}. {str(e)}') from e
+                    f'Failed to parse trigger config {trigger_config}. {str(e)}'
+                ) from e
             else:
                 print(f'Failed to parse trigger config for pipeline {pipeline_uuid}')
                 traceback.print_exc()
@@ -193,19 +283,26 @@ def build_triggers(
 def load_trigger_configs(
     content: str,
     pipeline_uuid: str = None,
+    repo_path: str = None,
     raise_exception: bool = False,
+    user=None,
 ) -> List[Trigger]:
-    yaml_config = yaml.safe_load(content) or {}
+    yaml_config = load_yaml(content) or {}
     trigger_configs = yaml_config.get('triggers') or {}
 
     return build_triggers(
-        trigger_configs, pipeline_uuid, raise_exception, repo_path=get_repo_path())
+        trigger_configs,
+        pipeline_uuid,
+        raise_exception,
+        repo_path=repo_path or get_repo_path(),
+    )
 
 
 def add_or_update_trigger_for_pipeline_and_persist(
     trigger: Trigger,
     pipeline_uuid: str,
     update_only_if_exists: bool = False,
+    old_trigger_name: str = None,
 ) -> Dict:
     trigger_configs_by_name = get_trigger_configs_by_name(pipeline_uuid)
 
@@ -214,13 +311,14 @@ def add_or_update_trigger_for_pipeline_and_persist(
     have, so we need to set "envs" on the updated trigger if it already exists.
     Otherwise, it will get overwritten when updating the trigger in code.
     """
-    existing_trigger = trigger_configs_by_name.get(trigger.name)
+    trigger_name = trigger.name if old_trigger_name is None else old_trigger_name
+    existing_trigger = trigger_configs_by_name.get(trigger_name)
     if existing_trigger is not None:
         trigger.envs = existing_trigger.get('envs', [])
     elif update_only_if_exists:
         return None
 
-    trigger_configs_by_name[trigger.name] = trigger.to_dict()
+    trigger_configs_by_name[trigger_name] = trigger.to_dict()
     yaml_config = dict(triggers=list(trigger_configs_by_name.values()))
     content = yaml.safe_dump(yaml_config)
     trigger_file_path = get_triggers_file_path(pipeline_uuid)