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

mage_ai/api/policies/BackfillPolicy.py

@@ -32,6 +32,7 @@ BackfillPolicy.allow_actions([
 
 BackfillPolicy.allow_read([
     'pipeline_run_dates',
+    'run_status_counts',
     'total_run_count',
 ] + BackfillPresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/policies/PipelinePolicy.py

@@ -150,6 +150,7 @@ PipelinePolicy.allow_query(
 PipelinePolicy.allow_query([
     'from_history_days',
     'include_schedules',
+    'repo_path',
     'search',
     'status[]',
     'tag[]',

mage_ai/api/policies/WorkspacePolicy.py

@@ -54,6 +54,7 @@ WorkspacePolicy.allow_write([
     'storage_access_mode',
     'storage_request_size',
     'task_definition',
+    'update_workspace_settings',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[

mage_ai/api/presenters/BackfillPresenter.py

@@ -49,6 +49,7 @@ class BackfillPresenter(BasePresenter):
         ):
             pipeline_run_dates = preview_run_dates(self.model)
             data['total_run_count'] = len(pipeline_run_dates)
+            data['run_status_counts'] = self.model.pipeline_run_status_counts
             if include_preview_runs:
                 start_idx = offset
                 end_idx = start_idx + limit

mage_ai/api/resources/GitBranchResource.py

@@ -1,16 +1,22 @@
 import os
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.data_preparation.git import REMOTE_NAME, Git, api
 from mage_ai.data_preparation.git.clients.base import Client as GitClient
-from mage_ai.data_preparation.git.utils import get_provider_from_remote_url
+from mage_ai.data_preparation.git.utils import (
+    get_oauth_access_token_for_user,
+    get_provider_from_remote_url,
+)
 from mage_ai.data_preparation.preferences import get_preferences
+from mage_ai.server.logger import Logger
 from mage_ai.shared.path_fixer import remove_base_repo_path
 from mage_ai.shared.strings import capitalize_remove_underscore_lower
 
+logger = Logger().new_server_logger(__name__)
+
 
 def build_file_object(obj):
     arr = []
@@ -53,7 +59,7 @@ class GitBranchResource(GenericResource):
             if remote_url:
                 provider = get_provider_from_remote_url(remote_url)
 
-            access_token = api.get_access_token_for_user(user, provider=provider)
+            access_token = get_oauth_access_token_for_user(user, provider=provider)
             if access_token:
                 branches = GitClient.get_client_for_provider(provider)(
                     access_token.token
@@ -64,7 +70,15 @@ class GitBranchResource(GenericResource):
             arr += [dict(name=branch) for branch in git_manager.branches]
 
             if include_remote_branches:
-                pass
+                try:
+                    git_manager.fetch()
+                    mage_remote = git_manager.origin
+                    arr += [
+                        dict(name=ref.name)
+                        for ref in mage_remote.refs
+                    ]
+                except Exception:
+                    logger.warning('Failed to fetch remote branches')
 
         return self.build_result_set(
             arr,
@@ -123,32 +137,27 @@ class GitBranchResource(GenericResource):
             **kwargs,
         )
 
-    async def update(self, payload, **kwargs):
-        query = kwargs.get('query') or {}
-
-        git_manager = self.get_git_manager(user=self.current_user)
-        action_type = payload.get('action_type')
-        action_payload = payload.get('action_payload', dict())
-        action_remote = action_payload.get('remote', None)
-        action_branch = action_payload.get('branch', None)
-
-        files = payload.get('files', None)
-        message = payload.get('message', None)
-
+    @classmethod
+    def get_oauth_config(
+        cls,
+        remote_url: str = None,
+        remote_name: str = None,
+        user=None,
+    ) -> Tuple[str, str, str, Dict]:
+        git_manager = cls.get_git_manager(user=user)
         url = None
-        remote_url = query.get('remote_url', None)
         if remote_url:
-            url = remote_url[0]
-        elif action_remote:
-            remote = git_manager.repo.remotes[action_remote]
+            url = remote_url
+        elif remote_name:
+            remote = git_manager.repo.remotes[remote_name]
             url = list(remote.urls)[0]
 
         provider = ProviderName.GITHUB
         if url:
             provider = get_provider_from_remote_url(url)
 
-        access_token = api.get_access_token_for_user(
-            self.current_user, provider=provider
+        access_token = get_oauth_access_token_for_user(
+            user, provider=provider
         )
         http_access_token = git_manager.get_access_token()
 
@@ -159,7 +168,7 @@ class GitBranchResource(GenericResource):
             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(
-                'email', self.current_user.email if self.current_user else None
+                'email', user.email if user else None
             )
             config_overwrite = dict(
                 username=user_from_api.get('username'),
@@ -168,6 +177,30 @@ class GitBranchResource(GenericResource):
         elif http_access_token:
             token = http_access_token
 
+        return token, provider, url, config_overwrite
+
+    async def update(self, payload, **kwargs):
+        query = kwargs.get('query') or {}
+
+        git_manager = self.get_git_manager(user=self.current_user)
+        action_type = payload.get('action_type')
+        action_payload = payload.get('action_payload', dict())
+        action_remote = action_payload.get('remote', None)
+        action_branch = action_payload.get('branch', None)
+
+        files = payload.get('files', None)
+        message = payload.get('message', None)
+
+        remote_url = query.get('remote_url', None)
+        if remote_url:
+            remote_url = remote_url[0]
+
+        token, provider, url, config_overwrite = self.get_oauth_config(
+            remote_url=remote_url,
+            remote_name=action_remote,
+            user=self.current_user,
+        )
+
         # Recreate git manager with updated config
         git_manager = self.get_git_manager(
             user=self.current_user, config_overwrite=config_overwrite

mage_ai/api/resources/GitCustomBranchResource.py

@@ -1,7 +1,7 @@
 from typing import Dict, List
 
 from mage_ai.api.resources.GitBranchResource import GitBranchResource
-from mage_ai.data_preparation.git import REMOTE_NAME, Git
+from mage_ai.data_preparation.git import REMOTE_NAME, Git, api
 from mage_ai.data_preparation.sync import AuthType
 
 
@@ -17,6 +17,34 @@ class GitCustomBranchResource(GitBranchResource):
             user=user,
         )
 
+    @classmethod
+    def create(cls, payload, user, **kwargs):
+        branch = payload.get('name')
+        remote = payload.get('remote')
+
+        token, _, url, config_overwrite = cls.get_oauth_config(
+            remote_name=remote,
+            user=user,
+        )
+
+        git_manager = cls.get_git_manager(
+            user=user, config_overwrite=config_overwrite
+        )
+
+        if remote:
+            api.switch_branch(
+                remote,
+                url,
+                branch,
+                token,
+                config_overwrite=config_overwrite,
+                user=user,
+            )
+        else:
+            git_manager.switch_branch(branch)
+
+        return cls(dict(name=git_manager.current_branch), user, **kwargs)
+
     @classmethod
     async def member(self, pk, user, **kwargs):
         resource = await GitBranchResource.member(pk, user, **kwargs)

mage_ai/api/resources/OauthResource.py

@@ -20,7 +20,7 @@ from mage_ai.authentication.oauth.utils import (
     refresh_token_for_client,
 )
 from mage_ai.authentication.providers.constants import NAME_TO_PROVIDER
-from mage_ai.data_preparation.git.api import get_oauth_client_id
+from mage_ai.data_preparation.git.utils import get_oauth_client_id
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
 

mage_ai/api/resources/PipelineResource.py

@@ -119,6 +119,12 @@ class PipelineResource(BaseResource):
         if include_schedules:
             include_schedules = include_schedules[0]
 
+        repo_path = query.get('repo_path', [None])
+        if repo_path:
+            repo_path = repo_path[0]
+        if not repo_path:
+            repo_path = get_repo_path(root_project=False)
+
         search_query = query.get('search', [None])
         if search_query:
             search_query = search_query[0]
@@ -174,11 +180,11 @@ class PipelineResource(BaseResource):
 
             if NO_TAGS_QUERY in tags:
                 pipeline_uuids_with_tags = set(cache.get_all_pipeline_uuids_with_tags())
-                all_pipeline_uuids = set(Pipeline.get_all_pipelines(get_repo_path()))
+                all_pipeline_uuids = set(Pipeline.get_all_pipelines(repo_path))
                 pipeline_uuids_without_tags = list(all_pipeline_uuids - pipeline_uuids_with_tags)
                 pipeline_uuids = pipeline_uuids + pipeline_uuids_without_tags
         else:
-            pipeline_uuids = Pipeline.get_all_pipelines(get_repo_path())
+            pipeline_uuids = Pipeline.get_all_pipelines(repo_path)
 
         if search_query:
             pipeline_uuids = list(filter(
@@ -206,7 +212,7 @@ class PipelineResource(BaseResource):
 
         async def get_pipeline(uuid: str) -> Pipeline:
             try:
-                return await Pipeline.get_async(uuid)
+                return await Pipeline.get_async(uuid, repo_path=repo_path)
             except Exception as err:
                 err_message = f'Error loading pipeline {uuid}: {err}.'
                 if err.__class__.__name__ == 'OSError' and 'Too many open files' in err.strerror:
@@ -217,7 +223,7 @@ class PipelineResource(BaseResource):
 
         def get_pipeline_with_config(uuid, config: Dict) -> Pipeline:
             try:
-                return Pipeline(uuid, config=config)
+                return Pipeline(uuid, config=config, repo_path=repo_path)
             except Exception as err:
                 err_message = f'Error loading pipeline sync {uuid}: {err}.'
                 if err.__class__.__name__ == 'OSError' and 'Too many open files' in err.strerror:
@@ -239,7 +245,7 @@ class PipelineResource(BaseResource):
                     ))
         else:
             for uuid in pipeline_uuids:
-                pipeline_dict = cache.get_model(dict(uuid=uuid))
+                pipeline_dict = cache.get_model(dict(uuid=uuid, repo_path=repo_path))
                 if pipeline_dict and pipeline_dict.get('pipeline'):
                     pipeline = get_pipeline_with_config(uuid, pipeline_dict['pipeline'])
                     if pipeline:

mage_ai/api/resources/PipelineRunResource.py

@@ -1,3 +1,4 @@
+from sqlalchemy import and_, desc, func
 from sqlalchemy.orm import selectinload
 
 from mage_ai.api.operations.constants import META_KEY_LIMIT, META_KEY_OFFSET
@@ -5,7 +6,10 @@ from mage_ai.api.resources.DatabaseResource import DatabaseResource
 from mage_ai.api.utils import get_query_timestamps
 from mage_ai.cache.tag import TagCache
 from mage_ai.data_preparation.models.block.utils import get_all_descendants
-from mage_ai.data_preparation.models.constants import PipelineType
+from mage_ai.data_preparation.models.constants import (
+    PIPELINE_RUN_STATUS_LAST_RUN_FAILED,
+    PipelineType,
+)
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import (
     ScheduleInterval,
@@ -94,9 +98,38 @@ class PipelineRunResource(DatabaseResource):
 
         repo_pipeline_schedule_ids = [s.id for s in PipelineSchedule.repo_query]
 
+        if status == PIPELINE_RUN_STATUS_LAST_RUN_FAILED:
+            """
+            The initial query below is used to get the last pipeline run retry
+            (or individual pipeline run if there are no retries) for each
+            grouping of pipeline runs with the same execution_date,
+            pipeline_uuid, and pipeline_schedule_id.
+            """
+            latest_pipeline_runs = PipelineRun.select(
+                PipelineRun.id,
+                func.row_number()
+                    .over(
+                        partition_by=(
+                            PipelineRun.execution_date,
+                            PipelineRun.pipeline_schedule_id,
+                            PipelineRun.pipeline_uuid,
+                        ),
+                        order_by=desc(PipelineRun.id))
+                    .label('row_number')
+            ).cte(name='latest_pipeline_runs')
+            query = (PipelineRun.select(
+                    PipelineRun,
+                )
+                .join(latest_pipeline_runs, and_(
+                    PipelineRun.id == latest_pipeline_runs.c.id,
+                    latest_pipeline_runs.c.row_number == 1,
+                ))
+            )
+        else:
+            query = PipelineRun.query
+
         results = (
-            PipelineRun
-            .query
+            query
             .filter(PipelineRun.pipeline_schedule_id.in_(repo_pipeline_schedule_ids))
             .options(selectinload(PipelineRun.block_runs))
             .options(selectinload(PipelineRun.pipeline_schedule))
@@ -120,8 +153,12 @@ class PipelineRunResource(DatabaseResource):
             results = results.filter(PipelineRun.pipeline_uuid.in_(pipeline_uuids))
         if pipeline_uuids_with_tags:
             results = results.filter(PipelineRun.pipeline_uuid.in_(pipeline_uuids_with_tags))
-        if status is not None:
+
+        if status == PIPELINE_RUN_STATUS_LAST_RUN_FAILED:
+            results = results.filter(PipelineRun.status == PipelineRun.PipelineRunStatus.FAILED)
+        elif status is not None:
             results = results.filter(PipelineRun.status == status)
+
         if statuses:
             results = results.filter(PipelineRun.status.in_(statuses))
         if start_timestamp is not None:

mage_ai/api/resources/PipelineScheduleResource.py

@@ -403,6 +403,8 @@ class PipelineScheduleResource(DatabaseResource):
         if updated_status == ScheduleStatus.ACTIVE and self.model.status == ScheduleStatus.INACTIVE:
             payload['last_enabled_at'] = datetime.now(tz=pytz.UTC)
 
+        old_name = self.model.name
+
         resource = super().update(payload)
         updated_model = resource.model
 
@@ -424,11 +426,13 @@ class PipelineScheduleResource(DatabaseResource):
             update_only_if_exists = (
                 not pipeline.should_save_trigger_in_code_automatically()
             )
+            old_trigger_name = old_name if old_name != updated_model.name else None
 
             add_or_update_trigger_for_pipeline_and_persist(
                 trigger,
                 pipeline.uuid,
                 update_only_if_exists=update_only_if_exists,
+                old_trigger_name=old_trigger_name,
             )
 
         return self

mage_ai/api/resources/PullRequestResource.py

@@ -3,9 +3,11 @@ from typing import Dict
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.authentication.oauth.constants import ProviderName
-from mage_ai.data_preparation.git import api
 from mage_ai.data_preparation.git.clients.base import Client as GitClient
-from mage_ai.data_preparation.git.utils import get_provider_from_remote_url
+from mage_ai.data_preparation.git.utils import (
+    get_oauth_access_token_for_user,
+    get_provider_from_remote_url,
+)
 
 
 def pull_request_to_dict(pr) -> Dict:
@@ -40,7 +42,7 @@ class PullRequestResource(GenericResource):
             if remote_url:
                 provider = get_provider_from_remote_url(remote_url)
 
-            access_token = api.get_access_token_for_user(user, provider=provider)
+            access_token = get_oauth_access_token_for_user(user, provider=provider)
 
             if access_token:
                 arr = GitClient.get_client_for_provider(provider)(
@@ -84,7 +86,7 @@ class PullRequestResource(GenericResource):
         remote_url = payload.get('remote_url')
         provider = get_provider_from_remote_url(remote_url)
 
-        access_token = api.get_access_token_for_user(user, provider=provider)
+        access_token = get_oauth_access_token_for_user(user, provider=provider)
         if not access_token:
             error.update(
                 dict(

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/api/resources/WorkspaceResource.py

@@ -123,14 +123,15 @@ class WorkspaceResource(GenericResource):
         error = ApiError.RESOURCE_ERROR.copy()
         try:
             action = payload.pop('action')
-            args = ignore_keys(payload, ['name', 'cluster_type'])
             if action == 'stop':
-                workspace.stop(**args)
+                workspace.stop(**payload)
             elif action == 'resume':
-                workspace.resume(**args)
+                workspace.resume(**payload)
+            elif action == 'patch':
+                workspace.update(payload)
             elif action == 'add_to_ingress':
                 if isinstance(workspace, KubernetesWorkspace):
-                    workspace.add_to_ingress(**args)
+                    workspace.add_to_ingress(**payload)
                 else:
                     raise Exception('This workspace does not support ingress.')
         except Exception as ex:

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.')