mage-ai 0.9.66__py3-none-any.whl → 0.9.68__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/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/decorators.py

@@ -103,3 +103,18 @@ def data_integration_query(function):
 
 def preprocesser_functions(function):
     return function
+
+
+def streaming_source(cls):
+    return cls
+
+
+def collect_decorated_objs(decorated_objs):
+    """
+    Method to collect the decorated objects (function or class)
+    """
+    def custom_code(obj):
+        decorated_objs.append(obj)
+        return obj
+
+    return custom_code

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
 

mage_ai/data_preparation/executors/streaming_pipeline_executor.py

@@ -10,7 +10,7 @@ from jinja2 import Template
 
 from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.logging.logger import DictLogger
-from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.shared.stream import StreamToLogger
 from mage_ai.data_preparation.shared.utils import get_template_vars
@@ -107,18 +107,28 @@ class StreamingPipelineExecutor(PipelineExecutor):
 
         if global_vars is None:
             global_vars = dict()
-        source_config = self.__interpolate_vars(
-            self.source_block.content,
-            global_vars=global_vars,
-        )
-        source = SourceFactory.get_source(
-            source_config,
-            checkpoint_path=os.path.join(
-                self.pipeline.pipeline_variables_dir,
-                'streaming_checkpoint',
-            ),
-        )
 
+        # Initialize source block
+        if self.source_block.language == BlockLanguage.PYTHON:
+            source = SourceFactory.get_python_source(
+                self.source_block.content,
+                global_vars=global_vars,
+            )
+        else:
+            # Default to YAML config
+            source_config = self.__interpolate_vars(
+                self.source_block.content,
+                global_vars=global_vars,
+            )
+            source = SourceFactory.get_source(
+                source_config,
+                checkpoint_path=os.path.join(
+                    self.pipeline.pipeline_variables_dir,
+                    'streaming_checkpoint',
+                ),
+            )
+
+        # Initialize destination blocks
         sinks_by_uuid = dict()
         for sink_block in self.sink_blocks:
             sinks_by_uuid[sink_block.uuid] = SinkFactory.get_sink(

mage_ai/data_preparation/git/api.py

@@ -15,8 +15,10 @@ from mage_ai.data_preparation.git.utils import (
 )
 from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, User
+from mage_ai.server.logger import Logger
 
 API_ENDPOINT = 'https://api.github.com'
+logger = Logger().new_server_logger(__name__)
 
 
 def get_oauth_client_id(provider: str) -> str:
@@ -262,4 +264,5 @@ def get_user(token: str, provider: str = ProviderName.GITHUB) -> Dict:
     try:
         return GitClient.get_client_for_provider(provider)(token).get_user()
     except Exception:
+        logger.exception('Error fetching user from git provider.')
         return dict()

mage_ai/data_preparation/git/clients/gitlab.py

@@ -3,14 +3,15 @@ from urllib.parse import quote_plus
 
 import requests
 
-from mage_ai.authentication.oauth.constants import GITLAB_HOST
 from mage_ai.data_preparation.git.clients.base import Client
+from mage_ai.settings import get_settings_value
+from mage_ai.settings.keys import GITLAB_HOST
 
 
 class GitlabClient(Client):
     def __init__(self, access_token: str):
         super().__init__(access_token)
-        self.hostname = GITLAB_HOST or 'https://gitlab.com'
+        self.hostname = get_settings_value(GITLAB_HOST, 'https://gitlab.com')
 
     def get_user(self):
         # For GitLab, we don't need to make a request to the API to get the username

mage_ai/data_preparation/logging/gcs_logger_manager.py

@@ -10,9 +10,9 @@ from mage_ai.shared.config import BaseConfig
 
 @dataclass
 class GCSConfig(BaseConfig):
-    path_to_credentials: str
     bucket: str
     prefix: str
+    path_to_credentials: str = None
 
 
 class GCSLoggerManager(LoggerManager):
@@ -23,9 +23,12 @@ class GCSLoggerManager(LoggerManager):
     ):
         super().__init__(repo_config=repo_config, **kwargs)
         self.gcs_config = GCSConfig.load(config=self.logging_config.destination_config)
-        credentials = service_account.Credentials.from_service_account_file(
-            self.gcs_config.path_to_credentials
-        )
+        if self.gcs_config.path_to_credentials:
+            credentials = service_account.Credentials.from_service_account_file(
+                self.gcs_config.path_to_credentials
+            )
+        else:
+            credentials = None
         self.gcs_client = storage.Client(credentials=credentials)
 
     def create_log_filepath_dir(self, path):