mage-ai 0.9.57__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/authentication/operation_history/utils.py

@@ -8,7 +8,7 @@ from mage_ai.authentication.operation_history.models import (
 )
 
 
-async def record_create_async(
+def record_create(
     resource_type: ResourceType,
     resource_uuid: str,
     user: str = None,
@@ -20,20 +20,20 @@ async def record_create_async(
         resource_uuid=resource_uuid,
         user=user,
     )
-    await reader.save_async(model)
+    reader.save(model)
 
     return model
 
 
-async def record_create_pipeline_async(resource_uuid: str, user: str = None) -> OperationHistory:
-    return await record_create_async(
+def record_create_pipeline(resource_uuid: str, user: str = None) -> OperationHistory:
+    return record_create(
         resource_type=ResourceType.PIPELINE.value,
         resource_uuid=resource_uuid,
         user=user,
     )
 
 
-async def record_detail_async(
+def record_detail(
     resource_type: ResourceType,
     resource_uuid: str,
     user: str = None,
@@ -45,20 +45,20 @@ async def record_detail_async(
         resource_uuid=resource_uuid,
         user=user,
     )
-    await reader.save_async(model)
+    reader.save(model)
 
     return model
 
 
-async def record_detail_pipeline_async(resource_uuid: str, user: str = None) -> OperationHistory:
-    return await record_detail_async(
+def record_detail_pipeline(resource_uuid: str, user: str = None) -> OperationHistory:
+    return record_detail(
         resource_type=ResourceType.PIPELINE.value,
         resource_uuid=resource_uuid,
         user=user,
     )
 
 
-async def load_records_async(
+def load_records(
     operation: OperationType = None,
     resource_type: ResourceType = None,
     timestamp_end: int = None,
@@ -66,7 +66,7 @@ async def load_records_async(
 ) -> List[OperationHistory]:
     reader = OperationHistoryReader()
 
-    return await reader.load_all_history_async(
+    return reader.load_all_history(
         operation=operation,
         resource_type=resource_type,
         timestamp_end=timestamp_end,
@@ -74,12 +74,12 @@ async def load_records_async(
     )
 
 
-async def load_pipelines_async(
+def load_pipelines(
     operation: OperationType = None,
     timestamp_end: int = None,
     timestamp_start: int = None,
 ) -> List[OperationHistory]:
-    return await load_records_async(
+    return load_records(
         operation=operation,
         resource_type=ResourceType.PIPELINE,
         timestamp_end=timestamp_end,
@@ -87,11 +87,11 @@ async def load_pipelines_async(
     )
 
 
-async def load_pipelines_detail_async(
+def load_pipelines_detail(
     timestamp_end: int = None,
     timestamp_start: int = None,
 ) -> List[OperationHistory]:
-    return await load_pipelines_async(
+    return load_pipelines(
         operation=OperationType.DETAIL,
         timestamp_end=timestamp_end,
         timestamp_start=timestamp_start,

mage_ai/authentication/permissions/constants.py

@@ -18,6 +18,7 @@ class EntityName(str, Enum):
     Chart = 'Chart'
     ClientPage = 'ClientPage'
     Cluster = 'Cluster'
+    CommandCenterItem = 'CommandCenterItem'
     ComputeCluster = 'ComputeCluster'
     ComputeConnection = 'ComputeConnection'
     ComputeService = 'ComputeService'
@@ -84,6 +85,10 @@ class EntityName(str, Enum):
     User = 'User'
     UserRole = 'UserRole'
     Variable = 'Variable'
+    VersionControlBranch = 'VersionControlBranch'
+    VersionControlFile = 'VersionControlFile'
+    VersionControlProject = 'VersionControlProject'
+    VersionControlRemote = 'VersionControlRemote'
     Widget = 'Widget'
     Workspace = 'Workspace'
 

mage_ai/authentication/providers/active_directory.py

@@ -8,7 +8,7 @@ import aiohttp
 from mage_ai.authentication.oauth.constants import (
     ACTIVE_DIRECTORY_CLIENT_ID as ACTIVE_DIRECTORY_MAGE_CLIENT_ID,
 )
-from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_ACTIVE_DIRECTORY
+from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.authentication.providers.oauth import OauthProvider
 from mage_ai.authentication.providers.sso import SsoProvider
 from mage_ai.authentication.providers.utils import get_base_url
@@ -20,7 +20,7 @@ from mage_ai.settings.sso import (
 
 
 class ADProvider(SsoProvider, OauthProvider):
-    provider = OAUTH_PROVIDER_ACTIVE_DIRECTORY
+    provider = ProviderName.ACTIVE_DIRECTORY
 
     def __init__(self):
         self.__validate()

mage_ai/authentication/providers/bitbucket.py

@@ -9,14 +9,14 @@ from mage_ai.authentication.oauth.constants import (
     BITBUCKET_HOST,
     BITBUCKET_OAUTH_KEY,
     BITBUCKET_OAUTH_SECRET,
-    OAUTH_PROVIDER_BITBUCKET,
+    ProviderName,
 )
 from mage_ai.authentication.providers.oauth import OauthProvider
 from mage_ai.authentication.providers.utils import get_base_url
 
 
 class BitbucketProvider(OauthProvider):
-    provider = OAUTH_PROVIDER_BITBUCKET
+    provider = ProviderName.BITBUCKET
 
     def __init__(self):
         self.hostname = BITBUCKET_HOST or 'https://bitbucket.org'

mage_ai/authentication/providers/constants.py

@@ -1,10 +1,4 @@
-from mage_ai.authentication.oauth.constants import (
-    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
-    OAUTH_PROVIDER_BITBUCKET,
-    OAUTH_PROVIDER_GHE,
-    OAUTH_PROVIDER_GOOGLE,
-    OAUTH_PROVIDER_OKTA,
-)
+from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.authentication.providers.active_directory import ADProvider
 from mage_ai.authentication.providers.bitbucket import BitbucketProvider
 from mage_ai.authentication.providers.ghe import GHEProvider
@@ -12,9 +6,9 @@ from mage_ai.authentication.providers.google import GoogleProvider
 from mage_ai.authentication.providers.okta import OktaProvider
 
 NAME_TO_PROVIDER = {
-    OAUTH_PROVIDER_ACTIVE_DIRECTORY: ADProvider,
-    OAUTH_PROVIDER_BITBUCKET: BitbucketProvider,
-    OAUTH_PROVIDER_GHE: GHEProvider,
-    OAUTH_PROVIDER_GOOGLE: GoogleProvider,
-    OAUTH_PROVIDER_OKTA: OktaProvider,
+    ProviderName.ACTIVE_DIRECTORY: ADProvider,
+    ProviderName.BITBUCKET: BitbucketProvider,
+    ProviderName.GHE: GHEProvider,
+    ProviderName.GOOGLE: GoogleProvider,
+    ProviderName.OKTA: OktaProvider,
 }

mage_ai/authentication/providers/ghe.py

@@ -8,7 +8,7 @@ import aiohttp
 from mage_ai.authentication.oauth.constants import (
     GHE_CLIENT_ID_ENV_VAR,
     GHE_CLIENT_SECRET_ENV_VAR,
-    OAUTH_PROVIDER_GHE,
+    ProviderName,
     get_ghe_hostname,
 )
 from mage_ai.authentication.providers.oauth import OauthProvider
@@ -16,7 +16,7 @@ from mage_ai.authentication.providers.utils import get_base_url
 
 
 class GHEProvider(OauthProvider):
-    provider = OAUTH_PROVIDER_GHE
+    provider = ProviderName.GHE
 
     def __init__(self):
         self.hostname = get_ghe_hostname()

mage_ai/authentication/providers/google.py

@@ -4,7 +4,7 @@ from typing import Awaitable, Dict
 
 import aiohttp
 
-from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_GOOGLE
+from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.authentication.providers.oauth import OauthProvider
 from mage_ai.authentication.providers.sso import SsoProvider
 from mage_ai.authentication.providers.utils import get_base_url
@@ -12,7 +12,7 @@ from mage_ai.settings.sso import GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET
 
 
 class GoogleProvider(SsoProvider, OauthProvider):
-    provider = OAUTH_PROVIDER_GOOGLE
+    provider = ProviderName.GOOGLE
 
     def __init__(self):
         if not GOOGLE_CLIENT_ID:

mage_ai/authentication/providers/okta.py

@@ -5,7 +5,7 @@ from typing import Awaitable, Dict
 import aiohttp
 from aiohttp import BasicAuth
 
-from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_OKTA
+from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.authentication.providers.oauth import OauthProvider
 from mage_ai.authentication.providers.sso import SsoProvider
 from mage_ai.authentication.providers.utils import get_base_url
@@ -13,7 +13,7 @@ from mage_ai.settings.sso import OKTA_CLIENT_ID, OKTA_CLIENT_SECRET, OKTA_DOMAIN
 
 
 class OktaProvider(SsoProvider, OauthProvider):
-    provider = OAUTH_PROVIDER_OKTA
+    provider = ProviderName.OKTA
 
     def __init__(self):
         self.hostname = OKTA_DOMAIN_URL

mage_ai/cache/base.py

@@ -13,7 +13,7 @@ from mage_ai.settings.repo import get_repo_path, get_variables_dir
 class BaseCache():
     cache_key = None
 
-    def __init__(self, repo_path: str = None, repo_config=None, root_project: bool = False):
+    def __init__(self, repo_path: str = None, repo_config=None, root_project: bool = True):
         self.root_project = root_project
         self.repo_path = repo_path or get_repo_path(root_project=self.root_project)
 
@@ -40,9 +40,12 @@ class BaseCache():
     def exists(self) -> bool:
         return self.get(self.cache_key) is not None
 
-    def get(self, key: str, refresh: bool = False, **kwargs) -> Union[Dict, List]:
+    def get(self, key: str = None, refresh: bool = False, **kwargs) -> Union[Dict, List]:
         if refresh or not self._temp_data:
-            self._temp_data = self.storage.read_json_file(self.build_path(key), None)
+            self._temp_data = self.storage.read_json_file(
+                self.build_path(key or self.cache_key),
+                None,
+            )
 
         if self._temp_data:
             return self._temp_data
@@ -72,6 +75,11 @@ class BaseCache():
 
         return os.path.join(dir_path, key)
 
+    def invalidate(self, key: str = None) -> None:
+        path = self.build_path(key or self.cache_key)
+        if self.storage.path_exists(path):
+            self.storage.remove(path)
+
     @property
     def file_path(self) -> str:
         return self.build_path(self.cache_key)

mage_ai/cache/block.py

@@ -3,9 +3,13 @@ import os
 from datetime import datetime
 from typing import Dict, List, Union
 
+import inflection
+
 from mage_ai.cache.base import BaseCache
 from mage_ai.cache.constants import CACHE_KEY_BLOCKS_TO_PIPELINE_MAPPING
-from mage_ai.cache.utils import build_pipeline_dict
+from mage_ai.cache.utils import build_block_dict, build_pipeline_dict
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.path_fixer import remove_base_repo_path_or_name
 
 
 class BlockCache(BaseCache):
@@ -14,16 +18,20 @@ class BlockCache(BaseCache):
     @classmethod
     async def initialize_cache(
         self,
-        replace: bool = False,
         caches: List[BaseCache] = None,
+        file_path: str = None,
+        replace: bool = False,
+        repo_path: str = None,
+        root_project: bool = True,
     ) -> 'BlockCache':
-        cache = self()
+        repo_path = repo_path or get_repo_path(root_project=root_project)
+        cache = self(repo_path=repo_path)
         if replace or not cache.exists():
-            await cache.initialize_cache_for_all_pipelines(caches=caches)
+            await cache.initialize_cache_for_all_pipelines(caches=caches, file_path=file_path)
 
         return cache
 
-    def build_key(self, block: Union[Dict]) -> str:
+    def build_key(self, block: Union[Dict], repo_path: str = None) -> str:
         """Generate cache key for block.
 
         Args:
@@ -35,27 +43,42 @@ class BlockCache(BaseCache):
         """
         block_type = ''
         block_uuid = ''
+        configuration = None
+        file_path = None
+        repo_path = repo_path or get_repo_path(root_project=False)
 
-        if type(block) is dict:
+        if isinstance(block, dict):
             block_type = block.get('type')
             block_uuid = block.get('uuid')
+            configuration = block.get('configuration') or {}
+            file_path = os.path.join(
+                repo_path,
+                inflection.pluralize(block_type),
+                block_uuid,
+            )
         else:
             block_type = block.type
             block_uuid = block.uuid
-        if not block_type or not block_uuid:
-            return None
-        return os.path.join(block_type, block_uuid)
+            configuration = block.configuration or {}
+            file_path = block.file_path
 
-    def get_pipelines(self, block) -> Dict:
-        pipelines_dict = {}
+        if configuration:
+            file_source = (configuration or {}).get('file_source') or {}
+            if file_source and (file_source or {}).get('path'):
+                return remove_base_repo_path_or_name((file_source or {}).get('path'))
+
+        return remove_base_repo_path_or_name(file_path)
+
+    def get_pipelines(self, block) -> List[Dict]:
+        pipelines = []
 
         mapping = self.get(self.cache_key)
         if mapping is not None:
             key = self.build_key(block)
             if key:
-                pipelines_dict = mapping.get(key, {})
+                return (mapping.get(key) or {}).get('pipelines', [])
 
-        return pipelines_dict
+        return pipelines
 
     def add_pipeline(self, block, pipeline) -> None:
         self.update_pipeline(block, pipeline, added_at=datetime.utcnow().timestamp())
@@ -73,8 +96,8 @@ class BlockCache(BaseCache):
         if mapping is None:
             mapping = {}
 
-        pipelines_dict = mapping.get(old_key, {})
-        mapping[new_key] = pipelines_dict
+        pipelines = mapping.get(old_key, [])
+        mapping[new_key] = pipelines
         mapping.pop(old_key, None)
 
         self.set(self.cache_key, mapping)
@@ -101,16 +124,45 @@ class BlockCache(BaseCache):
         if mapping is None:
             mapping = {}
 
-        pipelines_dict = mapping.get(key, {})
+        pipelines_arr = (mapping.get(key) or {}).get('pipelines', [])
 
         for pipeline in pipelines:
-            pipeline_uuid = pipeline.get('uuid') if type(pipeline) is dict else pipeline.uuid
-            pipelines_dict[pipeline_uuid] = build_pipeline_dict(
+            pipeline_uuid = None
+            repo_path = None
+
+            if isinstance(pipeline, dict):
+                pipeline_uuid = pipeline.get('uuid')
+                repo_path = pipeline.get('repo_path')
+            else:
+                pipeline_uuid = pipeline.uuid
+                repo_path = pipeline.repo_path
+
+            if not repo_path:
+                repo_path = get_repo_path(root_project=False)
+
+            repo_path = remove_base_repo_path_or_name(repo_path)
+
+            def __filter(
+                cache_dict: Dict,
+                pipeline_uuid=pipeline_uuid,
+                repo_path=repo_path,
+            ) -> bool:
+                pipeline_dict = cache_dict.get('pipeline') or {}
+
+                return pipeline_dict.get('uuid') != pipeline_uuid and \
+                    pipeline_dict.get('repo_path') != repo_path
+
+            pipelines_arr = list(filter(__filter, pipelines_arr))
+            pipelines_arr.append(build_pipeline_dict(
                 pipeline,
                 added_at=added_at,
-            )
+                repo_path=repo_path,
+            ))
 
-        mapping[key] = pipelines_dict
+        mapping[key] = dict(
+            block=build_block_dict(block),
+            pipelines=pipelines_arr,
+        )
 
         self.set(self.cache_key, mapping)
 
@@ -127,40 +179,72 @@ class BlockCache(BaseCache):
 
         self.set(self.cache_key, mapping)
 
-    def remove_pipeline(self, block, pipeline_uuid: str) -> None:
-        key = self.build_key(block)
+    def remove_pipeline(self, block, pipeline_uuid: str, repo_path: str = None) -> None:
+        repo_path = repo_path or get_repo_path(root_project=False)
+
+        key = self.build_key(block, repo_path=repo_path)
         if not key:
             return
 
         mapping = self.get(self.cache_key)
         if mapping is None:
-            mapping = {}
+            mapping = []
 
-        pipelines_dict = mapping.get(key, {})
-        pipelines_dict.pop(pipeline_uuid, None)
-        mapping[key] = pipelines_dict
+        repo_path = remove_base_repo_path_or_name(repo_path)
+
+        def __filter(cache_dict: Dict, pipeline_uuid=pipeline_uuid, repo_path=repo_path) -> bool:
+            pipeline_dict = cache_dict.get('pipeline') or {}
+            return pipeline_dict.get('uuid') != pipeline_uuid and \
+                pipeline_dict.get('repo_path') != repo_path
+
+        pipelines = (mapping.get(key) or {}).get('pipelines', [])
+
+        mapping[key] = dict(
+            block=build_block_dict(block),
+            pipelines=list(filter(__filter, pipelines)),
+        )
 
         self.set(self.cache_key, mapping)
 
-    async def initialize_cache_for_all_pipelines(self, caches: List[BaseCache] = None) -> None:
+    async def initialize_cache_for_all_pipelines(
+        self,
+        caches: List[BaseCache] = None,
+        file_path: str = None,
+    ) -> None:
         from mage_ai.data_preparation.models.pipeline import Pipeline
 
-        pipeline_uuids = Pipeline.get_all_pipelines(self.repo_path)
+        pipeline_uuids_and_repo_path = Pipeline.get_all_pipelines_all_projects(
+            self.repo_path,
+            include_repo_path=True,
+        )
 
         pipeline_dicts = await asyncio.gather(
-            *[Pipeline.load_metadata(uuid, raise_exception=False) for uuid in pipeline_uuids],
+            *[Pipeline.load_metadata(
+                uuid,
+                raise_exception=False,
+                repo_path=repo_path,
+            ) for uuid, repo_path in pipeline_uuids_and_repo_path],
         )
         pipeline_dicts = [p for p in pipeline_dicts if p is not None]
 
         mapping = {}
         for pipeline_dict in pipeline_dicts:
+            repo_path = pipeline_dict.get('repo_path')
+
             for block_dict in pipeline_dict.get('blocks', []):
-                key = self.build_key(block_dict)
+                key = self.build_key(block_dict, repo_path=repo_path)
                 if not key:
                     continue
                 if key not in mapping:
-                    mapping[key] = {}
-                mapping[key][pipeline_dict['uuid']] = build_pipeline_dict(pipeline_dict)
+                    mapping[key] = dict(
+                        block=build_block_dict(block_dict),
+                        pipelines=[],
+                    )
+
+                mapping[key]['pipelines'].append(build_pipeline_dict(
+                    pipeline_dict,
+                    repo_path=remove_base_repo_path_or_name(repo_path) if repo_path else repo_path,
+                ))
 
         if caches:
             for cache_class in caches:

mage_ai/cache/block_action_object/__init__.py

@@ -35,7 +35,8 @@ from mage_ai.data_preparation.templates.constants import (
 from mage_ai.data_preparation.templates.data_integrations.utils import (
     get_templates as get_templates_for_data_integrations,
 )
-from mage_ai.settings.repo import get_repo_path
+from mage_ai.settings.platform import build_repo_path_for_all_projects
+from mage_ai.shared.path_fixer import remove_base_repo_path
 from mage_ai.shared.strings import remove_extension_from_filename
 
 
@@ -44,7 +45,8 @@ def parse_block_file_absolute_path(block_file_absolute_path: str) -> Dict:
     file_directory_name = Block.file_directory_name(block_type)
 
     # This is the block_file_absolute_path without the repo_path
-    file_path = str(block_file_absolute_path).replace(get_repo_path(), '')
+    file_path = str(block_file_absolute_path)
+    file_path = remove_base_repo_path(str(block_file_absolute_path))
     if file_path.startswith(os.sep):
         file_path = file_path[1:]
 
@@ -151,7 +153,11 @@ class BlockActionObjectCache(BaseCache):
         if block:
             block_type = block.type
             file_path = os.path.join(block.file.dir_path, block.file.filename)
-            filename = os.path.join(*file_path.split(os.path.sep)[1:])
+            parts = file_path.split(os.path.sep)[1:]
+            if len(parts) == 0:
+                return
+
+            filename = os.path.join(*parts)
             language = block.language
             uuid = block.uuid
         else:
@@ -248,45 +254,48 @@ class BlockActionObjectCache(BaseCache):
                 include_content=True,
             )
 
-        for block_type in BlockType:
-            file_directory_name = Block.file_directory_name(block_type)
-            directory_full_path = os.path.join(get_repo_path(), file_directory_name)
-
-            for path in Path(directory_full_path).rglob('*'):
-                if not path.is_file():
-                    continue
-
-                block_file_absolute_path = path.absolute()
-                d = parse_block_file_absolute_path(block_file_absolute_path)
-
-                file_path = d.get('file_path')
-                filename = d.get('filename')
-                filename_parts = d.get('filename_parts')
-                language = d.get('language')
-
-                if '__init__.py' == filename:
-                    continue
-
-                if not language:
-                    continue
-
-                key = self.build_key_for_block_file(
-                    file_path,
-                    block_type,
-                    language,
-                    filename,
-                )
-
-                content = None
-                with open(block_file_absolute_path, 'r') as f:
-                    content = f.read()
-
-                mapping[OBJECT_TYPE_BLOCK_FILE][key] = dict(
-                    content=content,
-                    file_path=file_path,
-                    language=language,
-                    type=block_type,
-                    uuid='.'.join(filename_parts[:-1]),
-                )
+        for path_dict in build_repo_path_for_all_projects(mage_projects_only=True).values():
+            repo_path = path_dict['full_path']
+
+            for block_type in BlockType:
+                file_directory_name = Block.file_directory_name(block_type)
+                directory_full_path = os.path.join(repo_path, file_directory_name)
+
+                for path in Path(directory_full_path).rglob('*'):
+                    if not path.is_file():
+                        continue
+
+                    block_file_absolute_path = path.absolute()
+                    d = parse_block_file_absolute_path(block_file_absolute_path)
+
+                    file_path = d.get('file_path')
+                    filename = d.get('filename')
+                    filename_parts = d.get('filename_parts')
+                    language = d.get('language')
+
+                    if '__init__.py' == filename:
+                        continue
+
+                    if not language:
+                        continue
+
+                    key = self.build_key_for_block_file(
+                        file_path,
+                        block_type,
+                        language,
+                        filename,
+                    )
+
+                    content = None
+                    with open(block_file_absolute_path, 'r') as f:
+                        content = f.read()
+
+                    mapping[OBJECT_TYPE_BLOCK_FILE][key] = dict(
+                        content=content,
+                        file_path=file_path,
+                        language=language,
+                        type=block_type,
+                        uuid='.'.join(filename_parts[:-1]),
+                    )
 
         self.set(self.cache_key, mapping)

mage_ai/cache/constants.py

@@ -3,6 +3,7 @@ from enum import Enum
 CACHE_KEY_BLOCKS_TO_PIPELINE_MAPPING = 'blocks_to_pipeline_mapping'
 CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING = 'block_action_objects_mapping'
 CACHE_KEY_DBT_PROJECTS_PROFILES_MODELS = 'dbt_projects_profiles_models'
+CACHE_KEY_FILES = 'files'
 CACHE_KEY_PIPELINE_DETAILS_MAPPING = 'pipeline_details_mapping'
 CACHE_KEY_TAGS_TO_OBJECT_MAPPING = 'tags_to_object_mapping'