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

mage_ai/api/views.py

@@ -34,8 +34,7 @@ async def execute_operation(
     if request.error:
         return __render_error(handler, request.error, **tags)
 
-    action, options = __determine_action(
-        request, child=child, child_pk=child_pk, pk=pk)
+    action, options = __determine_action(request, child=child, child_pk=child_pk, pk=pk)
     try:
         base_operation = BaseOperation(
             action=action,
@@ -68,13 +67,15 @@ async def execute_operation(
     if error_response:
         return __render_error(handler, error_response, **tags)
 
-    info('Action: {} {} {} {} {}'.format(
-        action,
-        child or resource,
-        child_pk or pk,
-        resource if child else '',
-        pk if child else '',
-    ))
+    info(
+        'Action: {} {} {} {} {}'.format(
+            action,
+            child or resource,
+            child_pk or pk,
+            resource if child else '',
+            pk if child else '',
+        )
+    )
     api_time = end_time.timestamp() - start_time.timestamp()
     info(f'Latency: {api_time:.4f} seconds')
     tags = dict(
@@ -87,11 +88,13 @@ async def execute_operation(
     timing('api.time', api_time, tags)
     timing('sql.time', api_time, tags)
 
-    handler.write(simplejson.dumps(
-        response,
-        default=encode_complex,
-        ignore_nan=True,
-    ))
+    handler.write(
+        simplejson.dumps(
+            response,
+            default=encode_complex,
+            ignore_nan=True,
+        )
+    )
 
 
 def __determine_action(
@@ -148,19 +151,15 @@ def __meta(request) -> Dict:
 
 
 def __meta_keys(request) -> List[str]:
-    return list(
-        filter(
-            lambda x: x[0] == '_', [
-                k for k in request.query_arguments.keys()]))
+    return list(filter(lambda x: x[0] == '_', [k for k in request.query_arguments.keys()]))
 
 
 def __payload(request) -> Dict:
-    if 'Content-Type' in request.headers and \
-       'multipart/form-data' in request.headers.get('Content-Type'):
-
+    if 'Content-Type' in request.headers and 'multipart/form-data' in request.headers.get(
+        'Content-Type'
+    ):
         parts = request.body.decode('utf-8', 'ignore').split('\r\n')
-        idx = parts.index(
-            'Content-Disposition: form-data; name="json_root_body"')
+        idx = parts.index('Content-Disposition: form-data; name="json_root_body"')
         json_root_body = parts[idx + 2]
 
         return json.loads(json_root_body)
@@ -193,15 +192,18 @@ def __query(request) -> Dict:
         else:
             obj[key] = value
         return obj
+
     return reduce(_build, request.query_arguments.keys() - meta_keys, {})
 
 
 def __render_error(handler, error: Dict, **kwargs):
     __log_error(handler.request, error, **kwargs)
-    handler.write(dict(
-        error=error,
-        status=200,
-    ))
+    handler.write(
+        dict(
+            error=error,
+            status=200,
+        )
+    )
 
 
 def __log_error(
@@ -225,16 +227,21 @@ def __log_error(
             if child_pk:
                 endpoint += '/{}'.format(child_pk)
 
-    error('[{}] [ERROR] {} /{} [user: {}]: {}'.format(
-        datetime.utcnow(),
-        request.method,
-        endpoint,
-        user_id,
-        err,
-    ))
-
-    increment('api.error', tags=dict(
-        endpoint=endpoint,
-        error=err['type'] if isinstance(err, dict) else type(err).__name__,
-        resource=resource,
-    ))
+    error(
+        '[{}] [ERROR] {} /{} [user: {}]: {}'.format(
+            datetime.utcnow(),
+            request.method,
+            endpoint,
+            user_id,
+            err,
+        )
+    )
+
+    increment(
+        'api.error',
+        tags=dict(
+            endpoint=endpoint,
+            error=err['type'] if isinstance(err, dict) else type(err).__name__,
+            resource=resource,
+        ),
+    )

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']
 
@@ -294,7 +322,7 @@ class BlockActionObjectCache(BaseCache):
                     )
 
                     content = None
-                    with open(block_file_absolute_path, 'r') as f:
+                    with open(block_file_absolute_path, 'r', encoding='utf-8') as f:
                         content = f.read()
 
                     mapping[OBJECT_TYPE_BLOCK_FILE][key] = dict(

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