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

mage_ai/data_preparation/models/pipeline.py

@@ -15,8 +15,10 @@ import yaml
 from jinja2 import Template
 
 from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.cache.block import BlockCache
 from mage_ai.cache.pipeline import PipelineCache
 from mage_ai.data_preparation.models.block import Block, run_blocks, run_blocks_sync
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.block.data_integration.utils import (
     convert_outputs_to_data,
 )
@@ -60,6 +62,8 @@ from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.data_preparation.templates.utils import copy_template_directory
 from mage_ai.data_preparation.variable_manager import VariableManager
 from mage_ai.orchestration.constants import Entity
+from mage_ai.orchestration.notification.config import NotificationConfig
+from mage_ai.orchestration.notification.sender import NotificationSender
 from mage_ai.settings.platform import build_repo_path_for_all_projects
 from mage_ai.settings.platform.constants import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
@@ -82,6 +86,8 @@ class Pipeline:
         repo_config=None,
         catalog=None,
         use_repo_path: bool = False,
+        description: str = None,
+        tags: List[str] = None,
     ):
         self.block_configs = []
         self.blocks_by_uuid = {}
@@ -90,7 +96,7 @@ class Pipeline:
         self.concurrency_config = dict()
         self.created_at = None
         self.data_integration = None
-        self.description = None
+        self.description = description
         self.executor_config = dict()
         self.executor_type = None
         self.extensions = {}
@@ -101,7 +107,7 @@ class Pipeline:
         self.run_pipeline_in_one_process = False
         self.schedules = []
         self.settings = {}
-        self.tags = []
+        self.tags = tags or []
         self.type = PipelineType.PYTHON
         self.use_repo_path = use_repo_path
         self.uuid = uuid
@@ -215,7 +221,14 @@ class Pipeline:
             self.widget_configs
 
     @classmethod
-    def create(self, name, pipeline_type=PipelineType.PYTHON, repo_path=None):
+    def create(
+        self,
+        name: str,
+        description: str = None,
+        pipeline_type: PipelineType = PipelineType.PYTHON,
+        repo_path: str = None,
+        tags: List[str] = None,
+    ):
         """
         1. Create a new folder for pipeline
         2. Create a new yaml file to store pipeline config
@@ -231,14 +244,20 @@ class Pipeline:
         with open(os.path.join(pipeline_path, PIPELINE_CONFIG_FILE), 'w') as fp:
             yaml.dump(dict(
                 created_at=str(datetime.now(tz=pytz.UTC)),
+                description=description,
                 name=name,
+                tags=tags or [],
                 uuid=uuid,
                 type=format_enum(pipeline_type or PipelineType.PYTHON),
             ), fp)
+
         pipeline = Pipeline(
             uuid,
+            description=description,
             repo_path=repo_path,
+            tags=tags or [],
         )
+
         return pipeline
 
     @classmethod
@@ -322,7 +341,7 @@ class Pipeline:
         duplicate_pipeline_dict['name'] = duplicate_pipeline_uuid
         safe_write(
             duplicate_pipeline.config_path,
-            yaml.dump(duplicate_pipeline_dict)
+            yaml.dump(duplicate_pipeline_dict),
         )
 
         tags = duplicate_pipeline_dict.get('tags')
@@ -368,6 +387,54 @@ class Pipeline:
             IntegrationPipeline,
         )
 
+        config_path, repo_path = self._get_config_path(
+            uuid,
+            repo_path=repo_path,
+            all_projects=all_projects,
+            use_repo_path=use_repo_path,
+        )
+
+        if check_if_exists and not os.path.exists(config_path):
+            return None
+
+        pipeline = self(uuid, repo_path=repo_path, use_repo_path=use_repo_path)
+        if PipelineType.INTEGRATION == pipeline.type:
+            pipeline = IntegrationPipeline(uuid, repo_path=repo_path)
+
+        return pipeline
+
+    @classmethod
+    def get_config(
+        self,
+        uuid,
+        repo_path: str = None,
+        all_projects: bool = False,
+        use_repo_path: bool = False,
+    ):
+        config_path, _ = self._get_config_path(
+            uuid,
+            repo_path=repo_path,
+            all_projects=all_projects,
+            use_repo_path=use_repo_path,
+        )
+
+        metadata_path = os.path.join(config_path, PIPELINE_CONFIG_FILE)
+
+        if not os.path.exists(metadata_path):
+            return None
+
+        with open(metadata_path) as fp:
+            config = yaml.full_load(fp) or {}
+        return config
+
+    @classmethod
+    def _get_config_path(
+        self,
+        uuid,
+        repo_path: str = None,
+        all_projects: bool = False,
+        use_repo_path: bool = False,
+    ) -> Tuple[str, str]:
         if all_projects and not use_repo_path and project_platform_activated():
             from mage_ai.settings.platform.utils import get_pipeline_config_path
 
@@ -379,15 +446,7 @@ class Pipeline:
                 PIPELINES_FOLDER,
                 uuid,
             )
-
-        if check_if_exists and not os.path.exists(config_path):
-            return None
-
-        pipeline = self(uuid, repo_path=repo_path, use_repo_path=use_repo_path)
-        if PipelineType.INTEGRATION == pipeline.type:
-            pipeline = IntegrationPipeline(uuid, repo_path=repo_path)
-
-        return pipeline
+        return config_path, repo_path
 
     @classmethod
     async def load_metadata(
@@ -457,7 +516,7 @@ class Pipeline:
 
         if not os.path.exists(config_path):
             raise Exception(f'Pipeline {uuid} does not exist.')
-        async with aiofiles.open(config_path, mode='r') as f:
+        async with aiofiles.open(config_path, mode='r', encoding='utf-8') as f:
             config = yaml.safe_load(await f.read()) or {}
 
         if PipelineType.INTEGRATION == config.get('type'):
@@ -616,6 +675,7 @@ class Pipeline:
         analyze_outputs: bool = False,
         build_block_output_stdout: Callable[..., object] = None,
         global_vars=None,
+        retry_config=None,
         run_sensors: bool = True,
         run_tests: bool = True,
         update_status: bool = True,
@@ -632,6 +692,7 @@ class Pipeline:
             StreamingPipelineExecutor(self).execute(
                 build_block_output_stdout=build_block_output_stdout,
                 global_vars=global_vars,
+                retry_config=retry_config,
             )
         else:
             root_blocks = []
@@ -658,7 +719,7 @@ class Pipeline:
     def get_config_from_yaml(self):
         if not os.path.exists(self.config_path):
             raise Exception(f'Pipeline {self.uuid} does not exist in repo_path {self.repo_path}.')
-        with open(self.config_path) as fp:
+        with open(self.config_path, encoding='utf-8') as fp:
             config = yaml.full_load(fp) or {}
         return config
 
@@ -669,6 +730,16 @@ class Pipeline:
             config = json.load(f)
         return config
 
+    def get_notification_sender(self):
+        return NotificationSender(
+            NotificationConfig.load(
+                config=merge_dict(
+                    self.repo_config.notification_config,
+                    self.notification_config,
+                ),
+            ),
+        )
+
     def load_config_from_yaml(self):
         catalog = None
         if os.path.exists(self.catalog_config_path):
@@ -717,7 +788,8 @@ class Pipeline:
                 )
 
             language = c.get('language')
-            return Block.block_class_from_type(block_type, language=language, pipeline=self)(
+
+            return BlockFactory.block_class_from_type(block_type, language=language, pipeline=self)(
                 c.get('name'),
                 c.get('uuid'),
                 block_type,
@@ -931,6 +1003,8 @@ class Pipeline:
             include_outputs=include_outputs,
             sample_count=sample_count,
         )
+        if include_block_pipelines:
+            shared_kwargs['block_cache'] = BlockCache()
         blocks_data = await asyncio.gather(
             *[b.to_dict_async(**merge_dict(shared_kwargs, dict(
                 include_block_catalog=include_block_catalog,
@@ -991,6 +1065,7 @@ class Pipeline:
         old_uuid = None
         blocks_to_remove_from_cache = []
         block_uuids_to_add_to_cache = []
+        tags_to_remove_from_cache = []
         should_update_block_cache = False
         should_update_tag_cache = False
 
@@ -1043,6 +1118,10 @@ class Pipeline:
             old_tags = self.tags or []
 
             if sorted(new_tags) != sorted(old_tags):
+                tags_diff = set(old_tags).symmetric_difference(set(new_tags))
+                for tag in tags_diff:
+                    if tag in old_tags:
+                        tags_to_remove_from_cache.append(tag)
                 self.tags = new_tags
                 should_save = True
                 should_update_tag_cache = True
@@ -1229,7 +1308,9 @@ class Pipeline:
                         file_path = (configuration.get('file_source') or {}).get('path')
                         if file_path:
                             # Check for block name with period to avoid replacing a directory name
-                            new_file_path = file_path.replace(f'{block.name}.', f'{name}.')
+                            new_file_path = file_path.replace(
+                                f'{clean_name(block.name)}.', f'{clean_name(name)}.'
+                            )
                             configuration['file_source']['path'] = new_file_path
                             block_update_payload['configuration'] = configuration
                         blocks_to_remove_from_cache.append(block.to_dict())
@@ -1296,6 +1377,8 @@ class Pipeline:
 
             cache = await TagCache.initialize_cache()
 
+            for tag_uuid in tags_to_remove_from_cache:
+                cache.remove_pipeline(tag_uuid, self.uuid, self.repo_path)
             for tag_uuid in self.tags:
                 if old_uuid:
                     cache.remove_pipeline(tag_uuid, old_uuid, self.repo_path)
@@ -2065,7 +2148,7 @@ class Pipeline:
                 'Blocks cannot be added or removed when saving content, please try again.',
             )
 
-        content = yaml.dump(pipeline_dict)
+        content = yaml.dump(pipeline_dict, allow_unicode=True)
 
         safe_write(self.config_path, content)
 
@@ -2138,7 +2221,7 @@ class Pipeline:
                 'Blocks cannot be added or removed when saving content, please try again.',
             )
 
-        content = yaml.dump(pipeline_dict)
+        content = yaml.dump(pipeline_dict, allow_unicode=True)
 
         test_path = f'{self.config_path}.test'
         async with aiofiles.open(test_path, mode='w', encoding='utf-8') as fp:

mage_ai/data_preparation/models/utils.py

@@ -19,6 +19,12 @@ STRING_SERIALIZABLE_COLUMN_TYPES = set([
     'ObjectId',
 ])
 
+AMBIGUOUS_COLUMN_TYPES = set([
+    'mixed-integer',
+    'complex',
+    'unknown-array',
+])
+
 CAST_TYPE_COLUMN_TYPES = set([
     'Int64',
     'int64',

mage_ai/data_preparation/models/variable.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, List
 import numpy as np
 import pandas as pd
 import polars as pl
-from pandas.api.types import is_object_dtype
+from pandas.api.types import infer_dtype, is_object_dtype
 from pandas.core.indexes.range import RangeIndex
 
 from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
@@ -18,6 +18,7 @@ from mage_ai.data_preparation.models.constants import (
     VARIABLE_DIR,
 )
 from mage_ai.data_preparation.models.utils import (  # dask_from_pandas,
+    AMBIGUOUS_COLUMN_TYPES,
     STRING_SERIALIZABLE_COLUMN_TYPES,
     apply_transform_pandas,
     cast_column_types,
@@ -282,7 +283,9 @@ class Variable:
         else:
             self.__write_json(data)
 
-        self.write_metadata()
+        if self.variable_type != VariableType.SPARK_DATAFRAME:
+            # Not write json file in spark data directory to avoid read error
+            self.write_metadata()
 
     async def write_data_async(self, data: Any) -> None:
         """
@@ -315,7 +318,9 @@ class Variable:
         else:
             await self.__write_json_async(data)
 
-        self.write_metadata()
+        if self.variable_type != VariableType.SPARK_DATAFRAME:
+            # Not write json file in spark data directory to avoid read error
+            self.write_metadata()
 
     def write_metadata(self) -> None:
         """
@@ -578,10 +583,19 @@ class Variable:
             else:
                 series_non_null = df_col.dropna()
                 if len(series_non_null) > 0:
-                    coltype = type(series_non_null.iloc[0])
+                    sample_element = series_non_null.iloc[0]
+                    coltype = type(sample_element)
+                    coltype_inferred = infer_dtype(series_non_null)
                     if is_object_dtype(series_non_null.dtype):
                         if coltype.__name__ in STRING_SERIALIZABLE_COLUMN_TYPES:
                             cast_coltype = str
+                        # If the column is a "primitive" type, i.e. int/bool/etc and there is
+                        # a mix of types in the column, cast to string
+                        elif (
+                            not hasattr(sample_element, '__dict__')
+                            and coltype_inferred in AMBIGUOUS_COLUMN_TYPES
+                        ):
+                            cast_coltype = str
                         else:
                             cast_coltype = coltype
                         try:

mage_ai/data_preparation/repo_manager.py

@@ -310,8 +310,9 @@ def get_cluster_type(repo_path=None) -> Optional[ClusterType]:
 
 def set_project_uuid_from_metadata() -> None:
     global project_uuid
-    if os.path.exists(get_metadata_path()):
-        with open(get_metadata_path(), 'r', encoding='utf-8') as f:
+    metadata_path = get_metadata_path(root_project=True)
+    if os.path.exists(metadata_path):
+        with open(metadata_path, 'r', encoding='utf-8') as f:
             config = yml.load(f) or {}
             project_uuid = config.get('project_uuid')
 

mage_ai/data_preparation/shared/utils.py

@@ -25,7 +25,7 @@ def get_template_vars_no_db(include_python_libraries: Dict = None) -> Dict[str,
     try:
         from mage_ai.services.aws.secrets_manager.secrets_manager import get_secret
         kwargs['aws_secret_var'] = get_secret
-    except ImportError:
+    except Exception:
         pass
 
     try:

mage_ai/data_preparation/storage/gcs_storage.py

@@ -148,7 +148,7 @@ class GCSStorage(BaseStorage):
         df.write_parquet(buffer)
         buffer.seek(0)
         blob = self.bucket.blob(gcs_url_path(file_path))
-        blob.upload_from_string(buffer)
+        blob.upload_from_string(buffer.getvalue())
 
     @contextmanager
     def open_to_write(self, file_path: str) -> None:

mage_ai/data_preparation/templates/constants.py

@@ -608,6 +608,13 @@ TEMPLATES_ONLY_FOR_V2 = [
         name='MySQL',
         path='data_exporters/mysql.py',
     ),
+    dict(
+        block_type=BlockType.DATA_EXPORTER,
+        groups=[GROUP_DATABASES],
+        language=BlockLanguage.PYTHON,
+        name='OracleDB',
+        path='data_exporters/oracledb.py',
+    ),
     dict(
         block_type=BlockType.DATA_EXPORTER,
         groups=[GROUP_DATABASES],

mage_ai/data_preparation/templates/data_exporters/mysql.py

@@ -23,8 +23,8 @@ def export_data_to_mysql(df: DataFrame, **kwargs) -> None:
     with MySQL.with_config(ConfigFileLoader(config_path, config_profile)) as loader:
         loader.export(
             df,
-            None,
-            table_name,
+            schema_name=None,
+            table_name=table_name,
             index=False,  # Specifies whether to include index in exported table
             if_exists='replace',  # Specify resolution policy if table name already exists
         )

mage_ai/data_preparation/templates/data_exporters/oracledb.py

@@ -0,0 +1,27 @@
+from os import path
+
+from pandas import DataFrame
+
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.oracledb import OracleDB
+
+if 'data_exporter' not in globals():
+    from mage_ai.data_preparation.decorators import data_exporter
+
+
+@data_exporter
+def export_data_to_oracledb(df: DataFrame, **kwargs) -> None:
+    """
+    Template to export data to Oracledb.
+    """
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+    table_name = 'your_table_name'
+
+    with OracleDB.with_config(ConfigFileLoader(config_path, config_profile)) as exporter:
+        exporter.export(
+            df,
+            table_name=table_name,
+            if_exists='replace',  # Specify resolution policy if table name already exists
+        )

mage_ai/data_preparation/templates/repo/metadata.yaml

@@ -43,6 +43,7 @@ spark_config:
   # e.g. kwargs['context']['spark'] = spark_session
   custom_session_var_name: 'spark'
 
+help_improve_mage: true
 notification_config:
   alert_on:
     - trigger_failure