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

mage_ai/authentication/permissions/seed.py

@@ -1,5 +1,7 @@
+import asyncio
 import importlib
 import inspect
+import logging
 import os
 
 from mage_ai.api import policies
@@ -11,7 +13,7 @@ from mage_ai.authentication.permissions.constants import (
     PermissionAccess,
     PermissionCondition,
 )
-from mage_ai.orchestration.db import db_connection
+from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.models.oauth import Permission, Role, RolePermission, User
 
 KEY_ADMIN = 'admin'
@@ -32,6 +34,9 @@ ROLE_NAME_TO_KEY = {
 }
 
 
+logger = logging.getLogger(__name__)
+
+
 async def evaluate_condition(condition):
     if condition and inspect.isawaitable(condition):
         condition = await condition
@@ -224,6 +229,15 @@ async def attribute_rules(policy_class, rules):
     return mapping
 
 
+def bootstrap_permissions_sync(policy_names: str = None):
+    loop = asyncio.get_event_loop()
+    if loop is not None:
+        loop.create_task(bootstrap_permissions(policy_names))
+    else:
+        asyncio.run(bootstrap_permissions(policy_names))
+
+
+@safe_db_query
 async def bootstrap_permissions(policy_names: str = None):
     action_rules_mapping = {}
     query_rules_mapping = {}
@@ -248,7 +262,7 @@ async def bootstrap_permissions(policy_names: str = None):
         )
 
         model_name = policy_class.model_name()
-        print(f'Processing {model_name}...')
+        logger.info(f'Processing {model_name}...')
         action_rules_mapping[model_name] = await action_rules(policy_class)
         query_rules_mapping[model_name] = await attribute_rules(
             policy_class,

mage_ai/authentication/providers/oidc.py

@@ -1,3 +1,4 @@
+import json
 import urllib.parse
 import uuid
 from typing import Awaitable, Dict
@@ -11,7 +12,12 @@ from mage_ai.authentication.providers.sso import SsoProvider
 from mage_ai.authentication.providers.utils import get_base_url
 from mage_ai.server.logger import Logger
 from mage_ai.settings import get_settings_value
-from mage_ai.settings.keys import OIDC_CLIENT_ID, OIDC_CLIENT_SECRET, OIDC_DISCOVERY_URL
+from mage_ai.settings.keys import (
+    OIDC_CLIENT_ID,
+    OIDC_CLIENT_SECRET,
+    OIDC_DISCOVERY_URL,
+    OIDC_ROLES_MAPPING,
+)
 
 logger = Logger().new_server_logger(__name__)
 
@@ -25,6 +31,12 @@ class OidcProvider(OauthProvider, SsoProvider):
         self.client_secret = get_settings_value(OIDC_CLIENT_SECRET)
         self.__validate()
 
+        roles_mapping = get_settings_value(OIDC_ROLES_MAPPING)
+        if roles_mapping:
+            try:
+                self.roles_mapping = json.loads(roles_mapping)
+            except Exception:
+                logger.exception('Failed to parse OIDC roles mapping.')
         self.discover()
 
     def __validate(self):
@@ -115,6 +127,7 @@ class OidcProvider(OauthProvider, SsoProvider):
     async def get_user_info(self, access_token: str = None, **kwargs) -> Awaitable[Dict]:
         if access_token is None:
             raise Exception('Access token is required to fetch user info.')
+        mage_roles = []
         async with aiohttp.ClientSession() as session:
             async with session.get(
                 self.userinfo_endpoint,
@@ -127,8 +140,15 @@ class OidcProvider(OauthProvider, SsoProvider):
                 userinfo_resp = await response.json()
 
         email = userinfo_resp.get('email')
+        if hasattr(self, 'roles_mapping'):
+            for group in userinfo_resp.get('user_roles', []):
+                if group in self.roles_mapping:
+                    mage_roles.append(self.roles_mapping[group])
+        else:
+            mage_roles.extend(userinfo_resp.get('user_roles', []))
 
         return dict(
             email=email,
             username=userinfo_resp.get('preferred_username', email),
+            user_roles=mage_roles,
         )

mage_ai/autocomplete/utils.py

@@ -1,11 +1,12 @@
-from functools import reduce
-from mage_ai.shared.utils import files_in_path
-import aiofiles
 import importlib
 import os
 import pathlib
 import re
+from functools import reduce
 
+import aiofiles
+
+from mage_ai.shared.utils import files_in_path
 
 root_path = '/'.join(str(pathlib.Path(__file__).parent.resolve()).split('/')[:-2])
 
@@ -27,7 +28,7 @@ def add_file(acc, path):
 
     def __should_include(file_name):
         tup = os.path.splitext(file_name)
-        if (len(tup) >= 2):
+        if len(tup) >= 2:
             file_extension = tup[1]
             return file_extension in FILE_EXTENSIONS_TO_INCLUDE
 
@@ -60,6 +61,7 @@ def extract_all_imports(file_content, ignore_nesting=False):
         r'import [\w.]+',
         r'from [\w.]+ import [\w.]+ as [\w.]+',
         r'from [\w.]+ import [\w.]+',
+        r'from [\w.]+ import \(',
     ]
     regexes = []
 
@@ -89,7 +91,7 @@ async def build_file_content_mapping(paths, files):
         module_name = '.'.join(parts).replace('.py', '')
 
         if '__init__.py' == parts[-1]:
-            path_sub = '/'.join(parts[:len(parts) - 1])
+            path_sub = '/'.join(parts[: len(parts) - 1])
             files += [fn for fn in reduce(add_file, [path_sub], []) if fn != file_name]
             module_name = module_name.replace('.__init__', '')
 
@@ -101,10 +103,12 @@ async def build_file_content_mapping(paths, files):
                     importlib.import_module(module_name),
                     class_name,
                 )
-                methods_for_class[class_name] = list(filter(
-                    lambda x: not re.match('^_', x),
-                    dir(klass),
-                ))
+                methods_for_class[class_name] = list(
+                    filter(
+                        lambda x: not re.match('^_', x),
+                        dir(klass),
+                    )
+                )
             except ImportError as err:
                 print(err)
 

mage_ai/cache/base.py

@@ -69,7 +69,7 @@ class BaseCache():
 
     def build_path(self, key: str) -> str:
         dir_path = os.getenv(MAGE_CACHE_DIRECTORY_ENVIRONMENT_VARIABLE_NAME) or os.path.join(
-            get_variables_dir(repo_path=self.repo_path),
+            get_variables_dir(repo_path=self.repo_path, root_project=self.root_project),
             MAGE_CACHE_DIRECTORY_DEFAULT,
         )
 

mage_ai/cache/block.py

@@ -12,6 +12,7 @@ from mage_ai.data_preparation.models.constants import (
     BLOCK_LANGUAGE_TO_FILE_EXTENSION,
     BlockType,
 )
+from mage_ai.data_preparation.models.utils import warn_for_repo_path
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.path_fixer import remove_base_repo_path_or_name
 
@@ -50,7 +51,10 @@ class BlockCache(BaseCache):
         block_language = ''
         configuration = None
         file_path = None
-        repo_path = repo_path or get_repo_path(root_project=False)
+
+        warn_for_repo_path(repo_path)
+
+        repo_path = repo_path or get_repo_path()
 
         if isinstance(block, dict):
             block_type = block.get('type')
@@ -83,12 +87,12 @@ class BlockCache(BaseCache):
 
         return remove_base_repo_path_or_name(file_path)
 
-    def get_pipelines(self, block) -> List[Dict]:
+    def get_pipelines(self, block, repo_path: str) -> List[Dict]:
         pipelines = []
 
         mapping = self.get(self.cache_key)
         if mapping is not None:
-            key = self.build_key(block)
+            key = self.build_key(block, repo_path=repo_path)
             if key:
                 return (mapping.get(key) or {}).get('pipelines', [])
 
@@ -105,15 +109,15 @@ class BlockCache(BaseCache):
 
         return pipeline_count_mapping
 
-    def add_pipeline(self, block, pipeline) -> None:
-        self.update_pipeline(block, pipeline, added_at=datetime.utcnow().timestamp())
+    def add_pipeline(self, block, pipeline, repo_path: str) -> None:
+        self.update_pipeline(block, pipeline, repo_path, added_at=datetime.utcnow().timestamp())
 
-    def move_pipelines(self, new_block, old_block) -> None:
-        new_key = self.build_key(new_block)
+    def move_pipelines(self, new_block, old_block, repo_path: str) -> None:
+        new_key = self.build_key(new_block, repo_path=repo_path)
         if not new_key:
             return
 
-        old_key = self.build_key(old_block)
+        old_key = self.build_key(old_block, repo_path=repo_path)
         if not old_key:
             return
 
@@ -131,17 +135,19 @@ class BlockCache(BaseCache):
         self,
         block,
         pipeline,
+        repo_path: str,
         added_at: str = None,
     ) -> None:
-        self.update_pipelines(block, [pipeline], added_at)
+        self.update_pipelines(block, [pipeline], repo_path, added_at)
 
     def update_pipelines(
         self,
         block,
         pipelines,
+        repo_path: str,
         added_at: str = None,
     ) -> None:
-        key = self.build_key(block)
+        key = self.build_key(block, repo_path=repo_path)
         if not key:
             return
 
@@ -191,8 +197,8 @@ class BlockCache(BaseCache):
 
         self.set(self.cache_key, mapping)
 
-    def remove_block(self, block) -> None:
-        key = self.build_key(block)
+    def remove_block(self, block, repo_path: str) -> None:
+        key = self.build_key(block, repo_path=repo_path)
         if not key:
             return
 

mage_ai/cache/block_action_object/__init__.py

@@ -84,10 +84,14 @@ class BlockActionObjectCache(BaseCache):
     cache_key = CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING
 
     @classmethod
-    async def initialize_cache(self, replace: bool = False) -> 'BlockActionObjectCache':
+    async def initialize_cache(
+        self,
+        project: Project = None,
+        replace: bool = False,
+    ) -> 'BlockActionObjectCache':
         cache = self()
         if replace or not cache.exists():
-            await cache.initialize_cache_for_all_objects()
+            await cache.initialize_cache_for_all_objects(project=project)
 
         return cache
 
@@ -233,21 +237,43 @@ class BlockActionObjectCache(BaseCache):
 
         self.set(self.cache_key, mapping)
 
-    async def initialize_cache_for_all_objects(self) -> None:
+    async def initialize_cache_for_all_objects(self, project=None) -> None:
+        """
+        Initializes a cache for various types of objects across all projects.
+
+        This method populates a mapping dictionary with cached data for different types of objects,
+        including mage templates, custom block templates, and block files, from all projects.
+
+        - `OBJECT_TYPE_BLOCK_FILE`: Cache for block files.
+        - `OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE`: Cache for custom block templates.
+        - `OBJECT_TYPE_MAGE_TEMPLATE`: Cache for mage templates.
+
+        The caching process includes fetching templates relevant to all projects,
+        processing custom block templates from a directory,
+        and gathering data for block files from all projects' directories.
+
+        The cache is constructed and stored using the `set` method.
+
+        Returns:
+            None: This method does not return any value, it updates the cache for all objects.
+        """
         mapping = {
             OBJECT_TYPE_BLOCK_FILE: {},
             OBJECT_TYPE_CUSTOM_BLOCK_TEMPLATE: {},
             OBJECT_TYPE_MAGE_TEMPLATE: {},
         }
 
+        # Cache Mage templates
         mage_templates = TEMPLATES + TEMPLATES_ONLY_FOR_V2
-        if Project().is_feature_enabled(FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE):
+        project = project or Project()
+        if project.is_feature_enabled(FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE):
             mage_templates += get_templates_for_data_integrations()
 
         for block_action_object in mage_templates:
             key = self.build_key_for_mage_template(block_action_object)
             mapping[OBJECT_TYPE_MAGE_TEMPLATE][key] = block_action_object
 
+        # Cache custom block templates
         file_dicts = get_templates(DIRECTORY_FOR_BLOCK_TEMPLATES)
         file_dicts_flat = flatten_files(file_dicts)
         custom_block_templates = group_and_hydrate_files(file_dicts_flat, CustomBlockTemplate)
@@ -258,9 +284,11 @@ class BlockActionObjectCache(BaseCache):
                 include_content=True,
             )
 
+        # Cache block files
         paths_to_traverse = [dict(full_path=base_repo_path())]
         if project_platform_activated():
             paths_to_traverse = build_repo_path_for_all_projects(mage_projects_only=True).values()
+
         for path_dict in paths_to_traverse:
             repo_path = path_dict['full_path']
 

mage_ai/cache/file.py

@@ -12,7 +12,7 @@ from mage_ai.settings.utils import base_repo_path
 from mage_ai.shared.files import get_absolute_paths_from_all_files
 from mage_ai.shared.strings import size_of_string
 
-FILENAME_DELIMITER = '\t'
+FILENAME_DELIMITER = "\t"
 
 
 class FileCache(BaseCache):
@@ -23,7 +23,7 @@ class FileCache(BaseCache):
         self,
         absolute_paths: List[str],
         replace: bool = False,
-    ) -> 'FileCache':
+    ) -> "FileCache":
         repo_path = get_repo_path(root_project=True)
         cache = self(repo_path=repo_path)
 
@@ -33,7 +33,7 @@ class FileCache(BaseCache):
         return cache
 
     @classmethod
-    def initialize_cache_with_settings(self, replace: bool = False) -> 'FileCache':
+    def initialize_cache_with_settings(self, replace: bool = False) -> "FileCache":
         # Include all projects from platform
         # Relative to the base repo path
         base_path = base_repo_path()
@@ -45,10 +45,7 @@ class FileCache(BaseCache):
 
         # Include directories listed in command center settings and exclude those listed.
         settings = CommandCenterSettings.load_from_file_path()
-        if settings and \
-                settings.cache and \
-                settings.cache.files:
-
+        if settings and settings.cache and settings.cache.files:
             if settings.cache.files.directories:
                 paths.extend(settings.cache.files.directories.includes or [])
                 paths_excludes.extend(settings.cache.files.directories.excludes or [])
@@ -60,11 +57,12 @@ class FileCache(BaseCache):
         for project_name, path_dict in build_repo_path_for_all_projects(
             mage_projects_only=True,
         ).items():
-            path = path_dict.get('path')
-            if path and \
-                    project_name not in projects_excludes and \
-                    all([not path.startswith(ex) for ex in paths_excludes]):
-
+            path = path_dict.get("path")
+            if (
+                path
+                and project_name not in projects_excludes
+                and all([not path.startswith(ex) for ex in paths_excludes])
+            ):
                 paths.append(path)
 
         paths = [os.path.join(base_path, path) for path in paths]
@@ -89,17 +87,20 @@ class FileCache(BaseCache):
             try:
                 diff = Path(file_path).relative_to(base_path)
             except Exception as err:
-                print(f'[ERROR] FileCache.build_cache: {err}.')
+                print(f"[ERROR] FileCache.build_cache: {err}.")
                 diff = file_path
 
-            return f'{diff},{file_size},{file_modified_time}'
+            return f"{diff},{file_size},{file_modified_time}"
 
         arr_init = []
         for absolute_path in absolute_paths:
-            arr_init.extend(get_absolute_paths_from_all_files(
-                absolute_path,
-                parse_values=__parse_values,
-            ))
+            arr_init.extend(
+                get_absolute_paths_from_all_files(
+                    absolute_path,
+                    comparator=lambda absolute_path: not absolute_path.endswith(".pyc"),
+                    parse_values=__parse_values,
+                )
+            )
 
         if arr_init and len(arr_init) >= 1:
             arr_init.sort()
@@ -117,7 +118,9 @@ class FileCache(BaseCache):
 
             self.set(self.cache_key, FILENAME_DELIMITER.join(arr))
         else:
-            print('[WARNING] FileCache.build_cache: 0 files cached, this may be a mistake.')
+            print(
+                "[WARNING] FileCache.build_cache: 0 files cached, this may be a mistake."
+            )
 
         return arr
 

mage_ai/cache/pipeline.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, List, Union
 from mage_ai.cache.base import BaseCache
 from mage_ai.cache.constants import CACHE_KEY_PIPELINE_DETAILS_MAPPING
 from mage_ai.cache.utils import build_pipeline_dict, group_models_by_keys
+from mage_ai.data_preparation.models.utils import warn_for_repo_path
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.path_fixer import remove_base_repo_path_or_name
 
@@ -63,16 +64,17 @@ class PipelineCache(BaseCache):
             return pipeline_uuid
 
         if not repo_path:
-            repo_path = get_repo_path(root_project=False)
+            warn_for_repo_path(repo_path)
+            repo_path = get_repo_path()
 
         return ':'.join([remove_base_repo_path_or_name(repo_path), pipeline_uuid])
 
-    def get_model(self, model) -> Dict:
+    def get_model(self, model, repo_path: str = None) -> Dict:
         model_dict = {}
 
         mapping = self.get(self.cache_key)
         if mapping is not None:
-            key = self.build_key(model)
+            key = self.build_key(model, repo_path=repo_path)
             if key:
                 model_dict = mapping.get(key, {})
 
@@ -93,13 +95,14 @@ class PipelineCache(BaseCache):
         self,
         pipelines,
         added_at: str = None,
+        repo_path: str = None,
     ) -> None:
         mapping = self.get(self.cache_key)
         if mapping is None:
             mapping = {}
 
         for pipeline in pipelines:
-            key = self.build_key(pipeline)
+            key = self.build_key(pipeline, repo_path=repo_path)
             if not key:
                 continue
 
@@ -115,18 +118,21 @@ class PipelineCache(BaseCache):
         self,
         pipeline,
         added_at: str = None,
+        repo_path: str = None,
     ) -> None:
-        self.update_models([pipeline], added_at)
+        self.update_models([pipeline], added_at=added_at, repo_path=repo_path)
 
-    def add_model(self, model) -> None:
-        self.update_model(model, added_at=datetime.utcnow().timestamp())
+    def add_model(self, model, repo_path: str = None) -> None:
+        self.update_model(
+            model, added_at=datetime.utcnow().timestamp(), repo_path=repo_path
+        )
 
-    def move_model(self, new_model, old_model) -> None:
-        new_key = self.build_key(new_model)
+    def move_model(self, new_model, old_model, repo_path: str = None) -> None:
+        new_key = self.build_key(new_model, repo_path=repo_path)
         if not new_key:
             return
 
-        old_key = self.build_key(old_model)
+        old_key = self.build_key(old_model, repo_path=repo_path)
         if not old_key:
             return
 
@@ -140,8 +146,8 @@ class PipelineCache(BaseCache):
 
         self.set(self.cache_key, mapping)
 
-    def remove_model(self, model) -> None:
-        key = self.build_key(model)
+    def remove_model(self, model, repo_path: str = None) -> None:
+        key = self.build_key(model, repo_path=repo_path)
         if not key:
             return
 

mage_ai/cli/main.py

@@ -22,6 +22,7 @@ class OrderCommands(TyperGroup):
 
 app = typer.Typer(
     cls=OrderCommands,
+    pretty_exceptions_show_locals=False,
 )
 
 # Defaults

mage_ai/cluster_manager/aws/emr_cluster_manager.py

@@ -8,7 +8,6 @@ from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.services.aws.emr.config import EmrConfig
 from mage_ai.services.aws.emr.emr import describe_cluster, list_clusters
 from mage_ai.services.aws.emr.launcher import create_cluster
-from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.array import find
 from mage_ai.shared.hash import merge_dict
 
@@ -37,12 +36,12 @@ class EmrClusterManager(ClusterManager):
             is_active=c['Id'] == self.active_cluster_id,
         ) for c in valid_clusters]
 
-    def create_cluster(self, emr_config: Dict = None):
+    def create_cluster(self, repo_path: str, emr_config: Dict = None):
         emr_config = EmrConfig.load(
             config=merge_dict(get_repo_config().emr_config or dict(), emr_config or dict()))
 
         return create_cluster(
-            get_repo_path(),
+            repo_path,
             done_status=None,
             emr_config=emr_config,
             tags=dict(name=CLUSTER_NAME),
@@ -51,6 +50,7 @@ class EmrClusterManager(ClusterManager):
 
     def set_active_cluster(
         self,
+        repo_path: str = None,
         auto_creation: bool = True,
         auto_selection: bool = False,
         cluster_id=None,
@@ -61,13 +61,17 @@ class EmrClusterManager(ClusterManager):
             if len(clusters) > 0:
                 cluster_id = clusters[0]['id']
             elif auto_creation:
-                self.create_cluster(emr_config=emr_config)
+                self.create_cluster(repo_path, emr_config=emr_config)
         if cluster_id is None:
             return
 
         self.active_cluster_id = cluster_id
         emr_config = EmrConfig.load(
-            config=merge_dict(get_repo_config().emr_config or dict(), emr_config or dict()))
+            config=merge_dict(
+                get_repo_config(repo_path=repo_path).emr_config or dict(),
+                emr_config or dict(),
+            )
+        )
 
         # Fetch cluster master instance public DNS
         cluster_info = describe_cluster(cluster_id)

mage_ai/cluster_manager/config.py

@@ -65,7 +65,7 @@ class PostStart(BaseConfig):
         hook_path = config.get('hook_path')
         if hook_path and not os.path.isabs(hook_path):
             config['hook_path'] = os.path.join(
-                get_repo_path(), hook_path
+                get_repo_path(root_project=True), hook_path
             )
 
         return config
@@ -82,7 +82,7 @@ class LifecycleConfig(BaseConfig):
         pre_start_script_path = config.get('pre_start_script_path')
         if pre_start_script_path and not os.path.isabs(pre_start_script_path):
             config['pre_start_script_path'] = os.path.join(
-                get_repo_path(), pre_start_script_path
+                get_repo_path(root_project=True), pre_start_script_path
             )
 
         return config

mage_ai/cluster_manager/manage.py

@@ -67,7 +67,7 @@ def get_workspaces(cluster_type: ClusterType, **kwargs) -> List[Workspace]:
     """
     is_main_project = get_project_type() == ProjectType.MAIN
 
-    repo_path = get_repo_path()
+    repo_path = get_repo_path(root_project=True)
     projects_folder = os.path.join(repo_path, 'projects')
     if is_main_project:
         # sort by file last modified time

mage_ai/cluster_manager/workspace/base.py

@@ -33,7 +33,7 @@ class Workspace(abc.ABC):
 
     @classproperty
     def project_folder(self) -> str:
-        return os.path.join(get_repo_path(), 'projects')
+        return os.path.join(get_repo_path(root_project=True), 'projects')
 
     @property
     def config_path(self) -> str:

mage_ai/command_center/applications/factory.py

@@ -21,7 +21,7 @@ class ApplicationFactory(BaseFactory):
     async def fetch_items(self, **kwargs) -> List[Dict]:
         items = []
 
-        for item_dict in (self.page_history or []):
+        for item_dict in self.page_history or []:
             item_scored = self.filter_score(item_dict)
             if item_scored:
                 items.append(item_scored)
@@ -70,12 +70,15 @@ class ApplicationFactory(BaseFactory):
                 if app:
                     status = (app.get('state') or {}).get('status')
                     if ApplicationExpansionStatus.MINIMIZED == status:
-                        pass
-
-                    self.filter_score_mutate_accumulator(
-                        await build_close_application(uuid),
-                        items,
-                    )
+                        self.filter_score_mutate_accumulator(
+                            await build_open_application(uuid),
+                            items,
+                        )
+                    else:
+                        self.filter_score_mutate_accumulator(
+                            await build_close_application(uuid),
+                            items,
+                        )
 
             if not app:
                 self.filter_score_mutate_accumulator(