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

mage_ai/api/monitors/BaseMonitor.py

@@ -24,8 +24,6 @@ class BaseMonitor:
         if self.error.type:
             data['type'] = self.error.type
 
-        loop = asyncio.get_event_loop()
-
         kwargs = dict(
             resource=self.resource,
             **extract(
@@ -48,6 +46,7 @@ class BaseMonitor:
             ),
         )
 
+        loop = asyncio.get_event_loop()
         if loop is not None:
             loop.create_task(
                 UsageStatisticLogger().error(

mage_ai/api/policies/PipelinePolicy.py

@@ -85,9 +85,11 @@ PipelinePolicy.allow_write([
     'conditionals',
     'clone_pipeline_uuid',
     'custom_template_uuid',
+    'description',
     'extensions',
     'llm',
     'name',
+    'tags',
     'type',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/resources/BlockLayoutItemResource.py

@@ -6,7 +6,7 @@ import urllib.parse
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.widget.constants import ChartType
 from mage_ai.data_preparation.variable_manager import get_global_variables
@@ -73,7 +73,7 @@ class BlockLayoutItemResource(GenericResource):
                 data_source_type = data_source_config.get('type')
                 pipeline_uuid = data_source_config.get('pipeline_uuid')
 
-                block = Block.get_block(
+                block = BlockFactory.get_block(
                     block_config.get('name') or file_path or uuid,
                     block_uuid,
                     block_type,

mage_ai/api/resources/BlockResource.py

@@ -12,6 +12,7 @@ from mage_ai.cache.block_action_object.constants import (
     OBJECT_TYPE_MAGE_TEMPLATE,
 )
 from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.block.utils import clean_name
 from mage_ai.data_preparation.models.constants import (
     FILE_EXTENSION_TO_BLOCK_LANGUAGE,
@@ -163,7 +164,7 @@ class BlockResource(GenericResource):
             if payload.get('type') == BlockType.DBT and language == BlockLanguage.SQL and content:
                 from mage_ai.data_preparation.models.block.dbt import DBTBlock
 
-                dbt_block = DBTBlock(
+                dbt_block = DBTBlock.create(
                     name,
                     clean_name(name),
                     BlockType.DBT,
@@ -323,7 +324,7 @@ class BlockResource(GenericResource):
         if BlockType.DBT == block_type:
             from mage_ai.data_preparation.models.block.dbt import DBTBlock
 
-            block = DBTBlock(
+            block = DBTBlock.create(
                 block_uuid,
                 block_uuid,
                 block_type,
@@ -331,7 +332,7 @@ class BlockResource(GenericResource):
                 language=language,
             )
         else:
-            block = Block.get_block(block_uuid, block_uuid, block_type, language=language)
+            block = BlockFactory.get_block(block_uuid, block_uuid, block_type, language=language)
 
         if not block.exists():
             error.update(ApiError.RESOURCE_NOT_FOUND)

mage_ai/api/resources/CacheItemResource.py

@@ -50,7 +50,7 @@ class CacheItemResource(AsyncBaseResource):
         )
 
         if CacheItemType.DBT == item_type:
-            block = DBTBlock(
+            block = DBTBlock.create(
                 pk,
                 pk,
                 BlockType.DBT,

mage_ai/api/resources/GitBranchResource.py

@@ -152,14 +152,10 @@ class GitBranchResource(GenericResource):
         )
         http_access_token = git_manager.get_access_token()
 
+        config_overwrite = None
         token = None
         if access_token:
             token = access_token.token
-        elif http_access_token:
-            token = http_access_token
-
-        config_overwrite = None
-        if token:
             user_from_api = api.get_user(token, provider=provider)
             # Default to mage user email if no email is returned from API
             email = user_from_api.get(
@@ -169,6 +165,8 @@ class GitBranchResource(GenericResource):
                 username=user_from_api.get('username'),
                 email=email,
             )
+        elif http_access_token:
+            token = http_access_token
 
         # Recreate git manager with updated config
         git_manager = self.get_git_manager(

mage_ai/api/resources/PageBlockLayoutResource.py

@@ -3,7 +3,7 @@ import urllib.parse
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.presenters.block_layout.page import PageBlockLayout
 from mage_ai.shared.hash import extract, ignore_keys, merge_dict
 from mage_ai.shared.utils import clean_name
@@ -52,7 +52,7 @@ class PageBlockLayoutResource(GenericResource):
                     if len(parts) >= 2:
                         block_uuid_use = os.path.join(*parts[1:])
 
-                block = Block.get_block(
+                block = BlockFactory.get_block(
                     block_uuid_use,
                     block_uuid_use,
                     block_type,

mage_ai/api/resources/PipelineResource.py

@@ -362,6 +362,8 @@ class PipelineResource(BaseResource):
         clone_pipeline_uuid = payload.get('clone_pipeline_uuid')
         template_uuid = payload.get('custom_template_uuid')
         name = payload.get('name')
+        description = payload.get('description')
+        tags = payload.get('tags')
         pipeline_type = payload.get('type')
         llm_payload = payload.get('llm')
         pipeline = None
@@ -375,8 +377,10 @@ class PipelineResource(BaseResource):
         else:
             pipeline = Pipeline.create(
                 name,
+                description=description,
                 pipeline_type=pipeline_type,
                 repo_path=get_repo_path(),
+                tags=tags,
             )
 
             if llm_payload:
@@ -451,6 +455,15 @@ class PipelineResource(BaseResource):
             cache = await PipelineCache.initialize_cache()
             cache.add_model(resource.model)
 
+            tags = resource.model.tags
+            if tags:
+                from mage_ai.cache.tag import TagCache
+
+                cache = await TagCache.initialize_cache()
+
+                for tag_uuid in tags:
+                    cache.add_pipeline(tag_uuid, resource.model)
+
         self.on_create_callback = _on_create_callback
 
         return self(pipeline, user, **kwargs)

mage_ai/api/resources/PipelineRunResource.py

@@ -12,7 +12,7 @@ from mage_ai.data_preparation.models.triggers import (
     ScheduleStatus,
     ScheduleType,
 )
-from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.models.schedules import (
     BlockRun,
     PipelineRun,
@@ -241,6 +241,15 @@ class PipelineRunResource(DatabaseResource):
         has_next = results_size > limit
         final_end_idx = results_size - 1 if has_next else results_size
 
+        # Expire pipeline runs that are in progress so that the latest status is returned
+        # the next time they are fetched.
+        for run in results:
+            if run.status in (
+                PipelineRun.PipelineRunStatus.RUNNING,
+                PipelineRun.PipelineRunStatus.INITIAL,
+            ):
+                db_connection.session.expire(run)
+
         result_set = self.build_result_set(
             results[0:final_end_idx],
             user,

mage_ai/api/resources/SeedResource.py

@@ -6,6 +6,7 @@ class SeedResource(GenericResource):
     @classmethod
     async def create(self, payload, user, **kwargs):
         if payload.get('roles') and payload.get('permissions'):
-            await bootstrap_permissions()
+            policy_names = payload.get('policy_names')
+            await bootstrap_permissions(policy_names=policy_names)
 
         return self({}, user, **kwargs)

mage_ai/api/resources/SessionResource.py

@@ -11,9 +11,10 @@ from mage_ai.orchestration.db.models.oauth import Role, User
 from mage_ai.settings import (
     AUTHENTICATION_MODE,
     OAUTH_DEFAULT_ACCESS,
+    get_bool_value,
     get_settings_value,
 )
-from mage_ai.settings.keys import LDAP_DEFAULT_ACCESS
+from mage_ai.settings.keys import LDAP_DEFAULT_ACCESS, UPDATE_ROLES_ON_LOGIN
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
@@ -29,6 +30,9 @@ class SessionResource(BaseResource):
 
         oauth_client = kwargs.get('oauth_client')
 
+        update_roles_on_login = get_bool_value(
+            get_settings_value(UPDATE_ROLES_ON_LOGIN, default='False'))
+
         # Oauth sign in
         if token and provider:
             roles = []
@@ -57,6 +61,9 @@ class SessionResource(BaseResource):
                         email=email,
                         roles_new=roles,
                     )
+                elif update_roles_on_login and roles:
+                    user.update(roles_new=roles)
+
                 oauth_token = generate_access_token(user, oauth_client)
                 return self(oauth_token, user, **kwargs)
 
@@ -109,6 +116,11 @@ class SessionResource(BaseResource):
                     roles_new=roles,
                     username=email,
                 )
+            elif update_roles_on_login:
+                role_names = conn.get_user_roles(user_attributes)
+                if role_names:
+                    roles = Role.query.filter(Role.name.in_(role_names)).all()
+                    user.update(roles_new=roles)
 
             oauth_token = generate_access_token(user, kwargs['oauth_client'])
             return self(oauth_token, user, **kwargs)

mage_ai/authentication/permissions/constants.py

@@ -27,6 +27,7 @@ class EntityName(str, Enum):
     CustomTemplate = 'CustomTemplate'
     DataProvider = 'DataProvider'
     Database = 'Database'
+    Download = 'Download'
     EventMatcher = 'EventMatcher'
     EventRule = 'EventRule'
     ExecutionState = 'ExecutionState'
@@ -68,6 +69,7 @@ class EntityName(str, Enum):
     Scheduler = 'Scheduler'
     SearchResult = 'SearchResult'
     Secret = 'Secret'
+    Seed = 'Seed'
     Session = 'Session'
     SparkApplication = 'SparkApplication'
     SparkEnvironment = 'SparkEnvironment'

mage_ai/authentication/permissions/seed.py

@@ -171,13 +171,17 @@ async def action_rules(policy_class):
             if scope not in operations_mapping:
                 operations_mapping[scope] = {}
 
-            condition = config2.get('condition')
-            key = await get_key(policy_class, condition)
+            if not isinstance(config2, list):
+                config2 = [config2]
 
-            if key not in operations_mapping[scope]:
-                operations_mapping[scope][key] = []
+            for c in config2:
+                condition = c.get('condition')
+                key = await get_key(policy_class, condition)
+
+                if key not in operations_mapping[scope]:
+                    operations_mapping[scope][key] = []
 
-            operations_mapping[scope][key].append(operation)
+                operations_mapping[scope][key].append(operation)
 
     return operations_mapping
 
@@ -202,34 +206,40 @@ async def attribute_rules(policy_class, rules):
                 if config2 is None:
                     continue
 
-                condition = config2.get('condition') if config2 else None
-                key = await get_key(policy_class, condition)
+                if not isinstance(config2, list):
+                    config2 = [config2]
+
+                for c in config2:
+                    condition = c.get('condition')
+                    key = await get_key(policy_class, condition)
 
-                if key not in mapping[scope]:
-                    mapping[scope][key] = {}
+                    if key not in mapping[scope]:
+                        mapping[scope][key] = {}
 
-                if operation not in mapping[scope][key]:
-                    mapping[scope][key][operation] = []
+                    if operation not in mapping[scope][key]:
+                        mapping[scope][key][operation] = []
 
-                mapping[scope][key][operation].append(attribute)
+                    mapping[scope][key][operation].append(attribute)
 
     return mapping
 
 
-async def bootstrap_permissions():
+async def bootstrap_permissions(policy_names: str = None):
     action_rules_mapping = {}
     query_rules_mapping = {}
     read_rules_mapping = {}
     write_rules_mapping = {}
 
-    policy_names = []
-    for n in os.listdir(policies.__path__[0]):
-        if n.endswith('Policy.py') and n not in [
-            'BasePolicy.py',
-            'SeedPolicy.py',
-            'UserPolicy.py',
-        ]:
-            policy_names.append(n.replace('Policy.py', ''))
+    if policy_names is None:
+        policy_names = []
+        for n in os.listdir(policies.__path__[0]):
+            if n.endswith('Policy.py') and n not in [
+                'AsyncBasePolicy.py',
+                'BasePolicy.py',
+                'SeedPolicy.py',
+                'UserPolicy.py',
+            ]:
+                policy_names.append(n.replace('Policy.py', ''))
 
     for policy_name in policy_names:
         policy_class = getattr(
@@ -238,6 +248,7 @@ async def bootstrap_permissions():
         )
 
         model_name = policy_class.model_name()
+        print(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/active_directory.py

@@ -34,10 +34,11 @@ class ADProvider(SsoProvider, OauthProvider):
         self.client_secret = get_settings_value(ACTIVE_DIRECTORY_CLIENT_SECRET)
         self.__validate()
 
-        self.roles_mapping = get_settings_value(ACTIVE_DIRECTORY_ROLES_MAPPING)
-        if self.roles_mapping:
+        self.roles_mapping = {}
+        roles_mapping = get_settings_value(ACTIVE_DIRECTORY_ROLES_MAPPING)
+        if roles_mapping:
             try:
-                self.roles_mapping = json.loads(self.roles_mapping)
+                self.roles_mapping = json.loads(roles_mapping)
             except Exception:
                 logger.exception('Failed to parse roles mapping.')
 

mage_ai/authentication/providers/okta.py

@@ -1,104 +1,43 @@
-import urllib.parse
-import uuid
-from typing import Awaitable, Dict
-
-import aiohttp
-from aiohttp import BasicAuth
+from urllib.parse import unquote, urlparse
 
 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
+from mage_ai.authentication.providers.oidc import OidcProvider
 from mage_ai.settings import get_settings_value
 from mage_ai.settings.keys import OKTA_CLIENT_ID, OKTA_CLIENT_SECRET, OKTA_DOMAIN_URL
 
 
-class OktaProvider(SsoProvider, OauthProvider):
+class OktaProvider(
+    OidcProvider
+):  # Okta configuration uses OIDC so we can just subclass the OidcProvider
     provider = ProviderName.OKTA
 
     def __init__(self):
         self.hostname = get_settings_value(OKTA_DOMAIN_URL)
         self.client_id = get_settings_value(OKTA_CLIENT_ID)
         self.client_secret = get_settings_value(OKTA_CLIENT_SECRET)
+        self.parsed_url = urlparse(unquote(self.hostname))
+        if not self.parsed_url.scheme:
+            self.parsed_url = urlparse(unquote(f'https://{self.hostname}'))
         self.__validate()
 
-        if not self.hostname.startswith('https'):
-            self.hostname = f'https://{self.hostname}'
+        self.discovery_url = (
+            f'https://{self.parsed_url.netloc}/.well-known/openid-configuration'
+        )
+        self.discover()
 
     def __validate(self):
-        if not self.hostname:
-            raise Exception(
+        if not self.parsed_url.netloc:
+            raise ValueError(
                 'Okta hostname is empty. '
-                'Make sure the OKTA_DOMAIN_URL environment variable is set.')
+                'Make sure the OKTA_DOMAIN_URL environment variable is set.'
+            )
         if not self.client_id:
-            raise Exception(
+            raise ValueError(
                 'Okta client id is empty. '
-                'Make sure the OKTA_CLIENT_ID environment variable is set.')
+                'Make sure the OKTA_CLIENT_ID environment variable is set.'
+            )
         if not self.client_secret:
-            raise Exception(
+            raise ValueError(
                 'Okta client secret is empty. '
-                'Make sure the OKTA_CLIENT_SECRET environment variable is set.')
-
-    def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
-        base_url = get_base_url(redirect_uri)
-        redirect_uri_query = dict(
-            provider=self.provider,
-            redirect_uri=redirect_uri,
-        )
-        query = dict(
-            client_id=self.client_id,
-            redirect_uri=urllib.parse.quote_plus(
-                f'{base_url}/oauth',
-            ),
-            response_mode='query',
-            response_type='code',
-            scope='openid email profile',
-            state=uuid.uuid4().hex,
-        )
-        query_strings = []
-        for k, v in query.items():
-            query_strings.append(f'{k}={v}')
-
-        return dict(
-            url=f"{self.hostname}/oauth2/default/v1/authorize?{'&'.join(query_strings)}",
-            redirect_query_params=redirect_uri_query,
-        )
-
-    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
-        base_url = get_base_url(kwargs.get('redirect_uri'))
-        data = dict()
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                f'{self.hostname}/oauth2/default/v1/token',
-                headers={
-                    'Content-Type': 'application/x-www-form-urlencoded',
-                },
-                data=dict(
-                    grant_type='authorization_code',
-                    code=code,
-                    redirect_uri=f'{base_url}/oauth',
-                ),
-                auth=BasicAuth(self.client_id, self.client_secret),
-                timeout=20,
-            ) as response:
-                data = await response.json()
-
-        return data
-
-    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.')
-        async with aiohttp.ClientSession() as session:
-            async with session.get(
-                f'{self.hostname}/oauth2/default/v1/userinfo',
-                headers={
-                    'Authorization': f'Bearer {access_token}'
-                },
-                timeout=10,
-            ) as response:
-                userinfo_resp = await response.json()
-
-        return dict(
-            email=userinfo_resp.get('email'),
-            username=userinfo_resp.get('sub'),
-        )
+                'Make sure the OKTA_CLIENT_SECRET environment variable is set.'
+            )

mage_ai/cache/tag.py

@@ -144,6 +144,9 @@ class TagCache(BaseCache):
 
         mapping[key][KEY_FOR_PIPELINES] = pipelines_dict
 
+        if not pipelines_dict and mapping[key].keys() == set([KEY_FOR_PIPELINES]):
+            mapping.pop(key, None)
+
         self.set(self.cache_key, mapping)
 
     async def initialize_cache_for_all_objects(self) -> None:

mage_ai/cluster_manager/manage.py

@@ -70,11 +70,14 @@ def get_workspaces(cluster_type: ClusterType, **kwargs) -> List[Workspace]:
     repo_path = get_repo_path()
     projects_folder = os.path.join(repo_path, 'projects')
     if is_main_project:
+        # sort by file last modified time
         projects = [
-            f.name.split('.')[0]
+            f
             for f in os.scandir(projects_folder)
             if not f.is_dir() and f.name.endswith('.yaml')
         ]
+        projects.sort(key=lambda f: os.path.getmtime(f.path), reverse=True)
+        projects = [os.path.splitext(f.name)[0] for f in projects]
     else:
         instances = get_instances(cluster_type, **kwargs)
         projects = [instance.get('name') for instance in instances]

mage_ai/data_preparation/executors/k8s_block_executor.py

@@ -1,6 +1,9 @@
 from typing import Dict
 
+from jinja2 import Template
+
 from mage_ai.data_preparation.executors.block_executor import BlockExecutor
+from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.constants import DEFAULT_NAMESPACE
 from mage_ai.services.k8s.job_manager import JobManager as K8sJobManager
@@ -29,9 +32,13 @@ class K8sBlockExecutor(BlockExecutor):
             job_name_prefix = self.executor_config.job_name_prefix
 
         if self.executor_config.namespace:
-            namespace = self.executor_config.namespace
+            namespace = Template(self.executor_config.namespace).render(
+                variables=lambda x: global_vars.get(x) if global_vars else None,
+                **get_template_vars()
+            )
         else:
             namespace = DEFAULT_NAMESPACE
+
         job_manager = K8sJobManager(
             job_name=f'mage-{job_name_prefix}-block-{block_run_id}',
             logger=self.logger,

mage_ai/data_preparation/executors/k8s_pipeline_executor.py

@@ -1,8 +1,11 @@
 import traceback
 from typing import Dict
 
+from jinja2 import Template
+
 from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.services.k8s.config import K8sExecutorConfig
 from mage_ai.services.k8s.constants import DEFAULT_NAMESPACE
 from mage_ai.services.k8s.job_manager import JobManager as K8sJobManager
@@ -39,6 +42,7 @@ class K8sPipelineExecutor(PipelineExecutor):
     ) -> None:
         job_manager = self.get_job_manager(
             pipeline_run_id=pipeline_run_id,
+            global_vars=global_vars,
             **kwargs,
         )
         cmd = self._run_commands(
@@ -54,15 +58,22 @@ class K8sPipelineExecutor(PipelineExecutor):
     def get_job_manager(
         self,
         pipeline_run_id: int = None,
+        global_vars: Dict = None,
         **kwargs,
     ) -> K8sJobManager:
+        if global_vars is None:
+            global_vars = dict()
         if not self.executor_config.job_name_prefix:
             job_name_prefix = 'data-prep'
         else:
             job_name_prefix = self.executor_config.job_name_prefix
 
         if self.executor_config.namespace:
-            namespace = self.executor_config.namespace
+
+            namespace = Template(self.executor_config.namespace).render(
+                variables=lambda x: global_vars.get(x) if global_vars else None,
+                **get_template_vars()
+            )
         else:
             namespace = DEFAULT_NAMESPACE