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

mage_ai/api/resources/PipelineResource.py

@@ -19,9 +19,9 @@ from mage_ai.api.resources.BaseResource import BaseResource
 from mage_ai.api.resources.BlockResource import BlockResource
 from mage_ai.api.resources.LlmResource import LlmResource
 from mage_ai.authentication.operation_history.utils import (
-    load_pipelines_detail_async,
-    record_create_pipeline_async,
-    record_detail_pipeline_async,
+    load_pipelines_detail,
+    record_create_pipeline,
+    record_detail_pipeline,
 )
 from mage_ai.cache.pipeline import PipelineCache
 from mage_ai.data_preparation.models.constants import (
@@ -55,12 +55,41 @@ from mage_ai.server.kernels import PIPELINE_TO_KERNEL_NAME, KernelName
 from mage_ai.settings.platform import project_platform_activated
 from mage_ai.settings.platform.utils import get_pipeline_from_platform_async
 from mage_ai.settings.repo import get_repo_path
-from mage_ai.shared.array import find_index
+from mage_ai.shared.array import find, find_index
 from mage_ai.shared.hash import group_by, ignore_keys, merge_dict
 from mage_ai.shared.strings import is_number
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
+@safe_db_query
+def query_pipeline_schedules(pipeline_uuids: List[str]):
+    a = aliased(PipelineSchedule, name='a')
+    result = (
+        PipelineSchedule.
+        select(*[
+            a.created_at,
+            a.id,
+            a.name,
+            a.pipeline_uuid,
+            a.schedule_interval,
+            a.schedule_type,
+            a.start_time,
+            a.status,
+            a.updated_at,
+        ]).
+        filter(
+            a.pipeline_uuid.in_(pipeline_uuids),
+            or_(
+                a.repo_path.in_(Project().repo_path_for_database_query(
+                    'pipeline_schedules',
+                )),
+                a.repo_path.is_(None),
+            )
+        )
+    ).all()
+    return group_by(lambda x: x.pipeline_uuid, result)
+
+
 class PipelineResource(BaseResource):
     @classmethod
     @safe_db_query
@@ -122,7 +151,7 @@ class PipelineResource(BaseResource):
             timestamp_start = (datetime.utcnow() - timedelta(
                 hours=24 * int(from_history_days),
             )).timestamp()
-            history = await load_pipelines_detail_async(timestamp_start=timestamp_start)
+            history = load_pipelines_detail(timestamp_start=timestamp_start)
             history = sorted(
                 history,
                 key=lambda x: int(x.timestamp),
@@ -186,6 +215,17 @@ class PipelineResource(BaseResource):
                     print(err_message)
                     return None
 
+        def get_pipeline_with_config(uuid, config: Dict) -> Pipeline:
+            try:
+                return Pipeline(uuid, config=config)
+            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:
+                    raise Exception(err_message)
+                else:
+                    print(err_message)
+                    return None
+
         pipeline_uuids_miss = []
         pipelines = []
 
@@ -199,7 +239,9 @@ class PipelineResource(BaseResource):
             for uuid in pipeline_uuids:
                 pipeline_dict = cache.get_model(dict(uuid=uuid))
                 if pipeline_dict and pipeline_dict.get('pipeline'):
-                    pipelines.append(Pipeline(uuid, config=pipeline_dict['pipeline']))
+                    pipeline = get_pipeline_with_config(uuid, pipeline_dict['pipeline'])
+                    if pipeline:
+                        pipelines.append(pipeline)
                 else:
                     pipeline_uuids_miss.append(uuid)
 
@@ -210,34 +252,6 @@ class PipelineResource(BaseResource):
 
         pipelines = [p for p in pipelines if p is not None]
 
-        @safe_db_query
-        def query_pipeline_schedules(pipeline_uuids):
-            a = aliased(PipelineSchedule, name='a')
-            result = (
-                PipelineSchedule.
-                select(*[
-                    a.created_at,
-                    a.id,
-                    a.name,
-                    a.pipeline_uuid,
-                    a.schedule_interval,
-                    a.schedule_type,
-                    a.start_time,
-                    a.status,
-                    a.updated_at,
-                ]).
-                filter(
-                    a.pipeline_uuid.in_(pipeline_uuids),
-                    or_(
-                        a.repo_path.in_(Project().repo_path_for_database_query(
-                            'pipeline_schedules',
-                        )),
-                        a.repo_path.is_(None),
-                    )
-                )
-            ).all()
-            return group_by(lambda x: x.pipeline_uuid, result)
-
         mapping = {}
         if include_schedules:
             mapping = query_pipeline_schedules(pipeline_uuids)
@@ -282,7 +296,16 @@ class PipelineResource(BaseResource):
                         vals.append(val)
                         bools.append(val is None if not reverse_sort else val is not None)
                     elif 'triggers' == k.lower():
-                        val = len(p.blocks_by_uuid)
+                        val = len(p.schedules)
+                        vals.append(val)
+                        bools.append(val is None if not reverse_sort else val is not None)
+                    elif 'status' == k.lower():
+                        if len(p.schedules) == 0:
+                            val = 'no_schedules'
+                        elif find(lambda s: s.status == 'active', p.schedules) is not None:
+                            val = 'active'
+                        else:
+                            val = 'inactive'
                         vals.append(val)
                         bools.append(val is None if not reverse_sort else val is not None)
                     elif hasattr(p, k):
@@ -309,7 +332,7 @@ class PipelineResource(BaseResource):
             results = pipelines[start_index:end_index]
 
         results_size = len(results)
-        has_next = limit and total_count > limit
+        has_next = limit and results_size > limit
         final_end_idx = results_size - 1 if has_next else results_size
 
         arr = results[0:final_end_idx]
@@ -412,7 +435,7 @@ class PipelineResource(BaseResource):
                 template_uuid=template_uuid,
             )
 
-        await record_create_pipeline_async(
+        record_create_pipeline(
             resource_uuid=pipeline.uuid,
             user=user.id if user else None,
         )
@@ -464,7 +487,7 @@ class PipelineResource(BaseResource):
             if Project(pipeline.repo_config).is_feature_enabled(
                 FeatureUUID.OPERATION_HISTORY,
             ):
-                await record_detail_pipeline_async(
+                record_detail_pipeline(
                     resource_uuid=pipeline.uuid,
                     user=user.id if user else None,
                 )
@@ -478,6 +501,16 @@ class PipelineResource(BaseResource):
 
             await BlockCache.initialize_cache()
 
+        include_schedules = query.get('include_schedules', [False])
+        if include_schedules:
+            include_schedules = include_schedules[0]
+
+        if include_schedules:
+            mapping = query_pipeline_schedules([pipeline.uuid])
+
+            if mapping.get(pipeline.uuid):
+                pipeline.schedules = mapping[pipeline.uuid] or []
+
         return self(pipeline, user, **kwargs)
 
     @safe_db_query

mage_ai/api/resources/PipelineRunResource.py

@@ -129,6 +129,12 @@ class PipelineRunResource(DatabaseResource):
         if end_timestamp is not None:
             results = results.filter(PipelineRun.created_at <= end_timestamp)
 
+        limit = int((meta or {}).get(META_KEY_LIMIT, self.DEFAULT_LIMIT))
+
+        # No need to order the results if limit is 0
+        if limit == 0:
+            return results
+
         if order_by:
             arr = []
             for tup in order_by:
@@ -204,8 +210,7 @@ class PipelineRunResource(DatabaseResource):
         query arg and set it to True (e.g. in order to make the number of pipeline runs
         returned consistent across pages).
         """
-        if (meta or {}).get(META_KEY_LIMIT, None) is not None and \
-            total_results.count() >= 1 and \
+        if limit is not None and limit != 0 and total_count >= 1 and \
             not disable_retries_grouping and \
                 (pipeline_uuid is not None or pipeline_schedule_id is not None):
 

mage_ai/api/resources/ProjectResource.py

@@ -6,7 +6,6 @@ import yaml
 
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.cache.block_action_object import BlockActionObjectCache
-from mage_ai.data_preparation.models.block.dbt.profiles import PROFILES_FILE_NAME
 from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.data_preparation.repo_manager import (
@@ -113,7 +112,7 @@ class ProjectResource(GenericResource):
             )
             proc.check_returncode()
 
-            with open(os.path.join(directory, project_uuid, PROFILES_FILE_NAME), 'w') as f:
+            with open(os.path.join(directory, project_uuid, 'profiles.yml'), 'w') as f:
                 content = yaml.safe_dump({
                     project_uuid: dict(
                         outputs=dict(

mage_ai/api/resources/PullRequestResource.py

@@ -2,7 +2,7 @@ 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 OAUTH_PROVIDER_GITHUB
+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
@@ -36,7 +36,7 @@ class PullRequestResource(GenericResource):
         if repository:
             repository = repository[0]
 
-            provider = OAUTH_PROVIDER_GITHUB
+            provider = ProviderName.GITHUB
             if remote_url:
                 provider = get_provider_from_remote_url(remote_url)
 
@@ -56,12 +56,21 @@ class PullRequestResource(GenericResource):
         for key in [
             'base_branch',
             'compare_branch',
+            'body',
             'title',
         ]:
-            if key not in payload:
+            if key not in payload or payload.get(key) is None:
                 error.update(dict(message=f'Value for {key} is required but empty.'))
                 raise ApiError(error)
 
+        if payload.get('base_branch') == payload.get('compare_branch'):
+            error.update(
+                dict(
+                    message='Base branch and compare branch cannot be the same.',
+                )
+            )
+            raise ApiError(error)
+
         repository = payload.get('repository')
         if not repository:
             error.update(

mage_ai/api/resources/SearchResultResource.py

@@ -55,10 +55,10 @@ def filter_results(result: Dict) -> bool:
 class SearchResultResource(GenericResource):
     @classmethod
     async def create(self, payload: Dict, user, **kwargs):
-        pipeline_type = payload.get('pipeline_type', None)
-        query = payload.get('query', None)
-        ratio = payload.get('ratio', None)
-        search_type = payload.get('type', None)
+        pipeline_type = payload.get('pipeline_type', None) or None
+        query = payload.get('query', None) or None
+        ratio = payload.get('ratio', None) or None
+        search_type = payload.get('type', None) or None
 
         results = []
 

mage_ai/api/resources/VersionControlBranchResource.py

@@ -0,0 +1,85 @@
+import urllib.parse
+from typing import Dict
+
+from mage_ai.api.resources.AsyncBaseResource import AsyncBaseResource
+from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.version_control.models import Branch, Remote
+
+
+class VersionControlBranchResource(AsyncBaseResource):
+    @classmethod
+    async def collection(self, query: Dict, _meta: Dict, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+
+        remote = query.get('remote', [None])
+        if remote:
+            remote = remote[0]
+
+        if remote:
+            remote = Remote.load(name=remote)
+            remote.project = project
+            remote.hydrate()
+
+        model = Branch()
+        model.project = project
+        model.remote = remote
+
+        models = []
+        for m in [Branch.load(
+            name=line.strip(),
+        ) for line in model.list(include_all=True) if len(line) >= 1]:
+            m.project = project
+            m.remote = remote
+            models.append(m)
+
+        return self.build_result_set(
+            models,
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def create(self, payload: Dict, user: User, **kwargs):
+        model = self.hydrate_models(**kwargs)
+        model.name = payload.get('name')
+        model.create()
+
+        return self(model, user, **kwargs)
+
+    @classmethod
+    async def member(self, pk: str, user: User, **kwargs):
+        model = self.hydrate_models(**kwargs)
+        model.name = urllib.parse.unquote(pk)
+        return self(model, user, **kwargs)
+
+    async def update(self, payload: Dict, **kwargs):
+        self.model.update(**payload)
+
+    async def delete(self, **kwargs):
+        query = kwargs.get('query') or {}
+        force = query.get('force', [None])
+        if force:
+            force = force[0]
+
+        self.model.delete(force=force)
+
+    @classmethod
+    def hydrate_models(self, **kwargs):
+        project = kwargs.get('parent_model')
+
+        query = kwargs.get('query') or {}
+
+        remote = query.get('remote', [None])
+        if remote:
+            remote = remote[0]
+
+        if remote:
+            remote = Remote.load(name=remote)
+            remote.project = project
+            remote.hydrate()
+
+        model = Branch()
+        model.project = project
+        model.remote = remote
+
+        return model

mage_ai/api/resources/VersionControlFileResource.py

@@ -0,0 +1,65 @@
+import urllib.parse
+from typing import Dict
+
+from mage_ai.api.resources.AsyncBaseResource import AsyncBaseResource
+from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.version_control.models import File
+
+
+class VersionControlFileResource(AsyncBaseResource):
+    @classmethod
+    async def collection(self, query: Dict, _meta: Dict, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+
+        staged = query.get('staged', [None])
+        if staged:
+            staged = staged[0]
+
+        unstaged = query.get('unstaged', [None])
+        if unstaged:
+            unstaged = unstaged[0]
+
+        untracked = query.get('untracked', [None])
+        if untracked:
+            untracked = untracked[0]
+
+        model = File()
+        model.project = project
+
+        models = []
+        for m in [File.load(name=line) for line in model.list(
+            staged=staged,
+            unstaged=unstaged,
+            untracked=untracked,
+        ) if len(line) >= 1]:
+            m.project = project
+            models.append(m)
+
+        return self.build_result_set(
+            models,
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def create(self, payload: Dict, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+        model = File()
+        model.project = project
+        model.name = payload.get('name')
+        model.create()
+        return self(model, user, **kwargs)
+
+    @classmethod
+    async def member(self, pk: str, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+        model = File.load(name=urllib.parse.unquote(pk))
+        model.project = project
+        model.diff = model.detail()
+        return self(model, user, **kwargs)
+
+    async def update(self, payload: Dict, **kwargs):
+        self.model.update(**payload)
+
+    async def delete(self, **kwargs):
+        self.model.delete()

mage_ai/api/resources/VersionControlProjectResource.py

@@ -0,0 +1,46 @@
+import urllib.parse
+from typing import Dict
+
+from mage_ai.api.errors import ApiError
+from mage_ai.api.resources.AsyncBaseResource import AsyncBaseResource
+from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.shared.hash import ignore_keys
+from mage_ai.version_control.models import Project
+
+
+class VersionControlProjectResource(AsyncBaseResource):
+    @classmethod
+    async def collection(self, query: Dict, _meta: Dict, user: User, **kwargs):
+        return self.build_result_set(
+            Project.load_all(),
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def create(self, payload: Dict, user: User, **kwargs):
+        project = Project.create(payload['uuid'])
+        return self(project, user, **kwargs)
+
+    @classmethod
+    def get_model(self, pk, **kwargs):
+        model = Project.load(uuid=urllib.parse.unquote(pk))
+        if not model.exists():
+            raise ApiError(ApiError.RESOURCE_NOT_FOUND)
+
+        return model
+
+    @classmethod
+    async def member(self, pk: str, user: User, **kwargs):
+        return self(self.get_model(pk), user, **kwargs)
+
+    async def update(self, payload: Dict, **kwargs):
+        if 'email' in payload:
+            self.model.configure(email=payload.get('email'))
+        if 'name' in payload:
+            self.model.configure(name=payload.get('name'))
+
+        self.model.update(ignore_keys(payload, ['email', 'name']))
+
+    async def delete(self, **kwargs):
+        self.model.delete()

mage_ai/api/resources/VersionControlRemoteResource.py

@@ -0,0 +1,64 @@
+import urllib.parse
+from typing import Dict
+
+from mage_ai.api.resources.AsyncBaseResource import AsyncBaseResource
+from mage_ai.orchestration.db.models.oauth import User
+from mage_ai.shared.array import unique_by
+from mage_ai.version_control.models import Remote
+
+
+class VersionControlRemoteResource(AsyncBaseResource):
+    @classmethod
+    async def collection(self, query: Dict, _meta: Dict, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+        model = Remote()
+        model.project = project
+        lines = model.list()
+
+        models = []
+        for remote in unique_by(
+            [Remote.load_from_text(line) for line in lines if len(line) >= 1],
+            lambda x: f'{x.name}:{x.url}',
+        ):
+            remote.project = project
+            models.append(remote)
+
+        return self.build_result_set(
+            models,
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    async def create(self, payload: Dict, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+        model = Remote.load(**payload)
+        model.project = project
+        model.create()
+        return self(model, user, **kwargs)
+
+    @classmethod
+    async def member(self, pk: str, user: User, **kwargs):
+        project = kwargs.get('parent_model')
+
+        query = kwargs.get('query') or {}
+        url = query.get('url', [None])
+        if url:
+            url = url[0]
+
+        model = Remote.load(
+            name=urllib.parse.unquote(pk),
+            url=url,
+        )
+        model.project = project
+
+        return self(model, user, **kwargs)
+
+    async def update(self, payload: Dict, **kwargs):
+        if payload.get('fetch'):
+            self.model.update(fetch=True)
+        if payload.get('set_url'):
+            self.model.update(set_url=True)
+
+    async def delete(self, **kwargs):
+        self.model.delete()

mage_ai/authentication/oauth/constants.py

@@ -1,4 +1,5 @@
 import os
+from enum import Enum
 from typing import Optional
 
 ACTIVE_DIRECTORY_CLIENT_ID = '51aec820-9d49-40a9-b046-17c1f28f620d'
@@ -6,12 +7,17 @@ ACTIVE_DIRECTORY_CLIENT_ID = '51aec820-9d49-40a9-b046-17c1f28f620d'
 GITHUB_CLIENT_ID = '8577f13ddc81e2848b07'
 GITHUB_STATE = '1337'
 
-OAUTH_PROVIDER_ACTIVE_DIRECTORY = 'active_directory'
-OAUTH_PROVIDER_BITBUCKET = 'bitbucket'
-OAUTH_PROVIDER_GITHUB = 'github'
-OAUTH_PROVIDER_GHE = 'ghe'
-OAUTH_PROVIDER_GOOGLE = 'google'
-OAUTH_PROVIDER_OKTA = 'okta'
+
+class ProviderName(str, Enum):
+    ACTIVE_DIRECTORY = 'active_directory'
+    BITBUCKET = 'bitbucket'
+    GITHUB = 'github'
+    GHE = 'ghe'
+    GOOGLE = 'google'
+    OKTA = 'okta'
+
+
+VALID_OAUTH_PROVIDERS = [e.value for e in ProviderName]
 
 GHE_CLIENT_ID_ENV_VAR = 'GHE_CLIENT_ID'
 GHE_CLIENT_SECRET_ENV_VAR = 'GHE_CLIENT_SECRET'
@@ -24,16 +30,7 @@ BITBUCKET_OAUTH_SECRET = os.getenv('BITBUCKET_OAUTH_SECRET')
 DEFAULT_GITHUB_HOSTNAME = 'https://github.com'
 
 GIT_OAUTH_PROVIDERS = [
-    OAUTH_PROVIDER_BITBUCKET,
-]
-
-VALID_OAUTH_PROVIDERS = [
-    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
-    OAUTH_PROVIDER_BITBUCKET,
-    OAUTH_PROVIDER_GHE,
-    OAUTH_PROVIDER_GITHUB,
-    OAUTH_PROVIDER_GOOGLE,
-    OAUTH_PROVIDER_OKTA,
+    ProviderName.BITBUCKET,
 ]
 
 

mage_ai/authentication/operation_history/models.py

@@ -1,7 +1,5 @@
-import asyncio
 import json
 import os
-from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from typing import Dict, List, Union
@@ -14,9 +12,9 @@ from mage_ai.authentication.operation_history.constants import (
 )
 from mage_ai.data_preparation.repo_manager import RepoConfig, get_repo_config
 from mage_ai.data_preparation.storage.local_storage import LocalStorage
+from mage_ai.settings.platform.constants import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.environments import is_debug
-from mage_ai.shared.io import read_last_line_async
 
 
 @dataclass
@@ -49,10 +47,14 @@ class OperationHistory:
 
 class OperationHistoryReader:
     def __init__(self, repo_config=None, repo_path: str = None):
-        self.repo_path = repo_path or get_repo_path()
+        self.root_project = project_platform_activated()
+        self.repo_path = repo_path or get_repo_path(root_project=self.root_project)
 
         if repo_config is None:
-            self.repo_config = get_repo_config(repo_path=self.repo_path)
+            self.repo_config = get_repo_config(
+                repo_path=self.repo_path,
+                root_project=self.root_project,
+            )
         elif isinstance(repo_config, dict):
             self.repo_config = RepoConfig.from_dict(repo_config)
         else:
@@ -102,17 +104,7 @@ class OperationHistoryReader:
 
         return os.path.join(dir_path, filename)
 
-    async def load_most_recent_history(self, timestamp: int = None) -> OperationHistory:
-        try:
-            file_path = self.build_file_path()
-            if file_path and self.storage.path_exists(file_path):
-                line = await read_last_line_async()
-                return OperationHistory(**json.loads(line))
-        except Exception as err:
-            if is_debug():
-                print(f'[ERROR] OperationHistoryReader.load_most_recent_history: {err}')
-
-    async def load_all_history_async(
+    def load_all_history(
         self,
         operation: OperationType = None,
         resource_type: ResourceType = None,
@@ -150,7 +142,7 @@ class OperationHistoryReader:
             if not self.storage.path_exists(file_path):
                 continue
 
-            text = await self.storage.read_async(file_path)
+            text = self.storage.read(file_path)
 
             if text:
                 for line in text.split('\n'):
@@ -166,26 +158,13 @@ class OperationHistoryReader:
                             arr.append(record)
                     except json.decoder.JSONDecodeError as err:
                         if is_debug():
-                            print(f'[ERROR] OperationHistoryReader.load_all_history_async: {err}')
+                            print(f'[ERROR] OperationHistoryReader.load_all_history: {err}')
 
         return arr
 
-    async def save(self, operation_history: OperationHistory) -> None:
+    def save(self, operation_history: OperationHistory) -> None:
         text = json.dumps(operation_history.to_dict())
         file_path = self.build_file_path(timestamp=operation_history.timestamp)
         self.storage.makedirs(os.path.dirname(file_path), exist_ok=True)
         with self.storage.open_to_write(file_path, append=True) as f:
             f.write(f'{text}\n')
-
-    async def save_async(self, operation_history: OperationHistory) -> None:
-        last_record = await self.load_most_recent_history()
-
-        if not last_record or last_record != operation_history:
-            try:
-                asyncio.get_running_loop()
-                with ThreadPoolExecutor(1) as pool:
-                    pool.submit(lambda: asyncio.run(self.save(operation_history))).result()
-            except Exception as err:
-                if is_debug():
-                    print(f'[ERROR] OperationHistoryReader.save_async: {err}')
-                asyncio.run(self.save(operation_history))