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

mage_ai/data_preparation/models/block/global_data_product/__init__.py

@@ -2,6 +2,7 @@ from logging import Logger
 from typing import Dict, List
 
 from mage_ai.data_preparation.models.block import Block
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 
 
 class GlobalDataProductBlock(Block):
@@ -37,13 +38,35 @@ class GlobalDataProductBlock(Block):
         # Avoid circular dependency of Pipeline class
         from mage_ai.orchestration.triggers import global_data_product as trigger
 
-        trigger.trigger_and_check_status(
-            self.get_global_data_product(),
+        global_data_product = self.get_global_data_product()
+
+        pipeline_run = trigger.trigger_and_check_status(
+            global_data_product,
             block=self,
             from_notebook=from_notebook,
             logger=logger,
             logging_tags=logging_tags,
+            poll_interval=30,
+            remote_blocks=global_vars.get('remote_blocks'),
             variables=global_vars.get('variables') if global_vars else None,
         )
 
+        if pipeline_run and pipeline_run.pipeline:
+            pipeline = pipeline_run.pipeline
+
+            arr = []
+
+            for block in global_data_product.get_blocks():
+                if block and block.uuid:
+                    arr.append(RemoteBlock.load(
+                        block_uuid=block.uuid,
+                        execution_partition=pipeline_run.execution_partition,
+                        pipeline_uuid=pipeline.uuid,
+                        repo_path=pipeline.repo_path,
+                    ))
+
+            return [
+                dict(remote_blocks=arr),
+            ]
+
         return []

mage_ai/data_preparation/models/block/integration/__init__.py

@@ -155,7 +155,7 @@ class IntegrationBlock(Block):
                 proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
 
                 for line in proc.stdout:
-                    f.write(line.decode()),
+                    f.write(line.decode())
                     print_log_from_line(
                         line,
                         config=config,

mage_ai/data_preparation/models/block/remote/models.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from typing import Any, List
+
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class RemoteBlock(BaseDataClass):
+    block_uuid: str
+    execution_partition: str = None
+    pipeline_uuid: str = None
+    repo_path: str = None
+
+    def __post_init__(self):
+        self._block = None
+        self._pipeline = None
+
+    @property
+    def pipeline(self):
+        from mage_ai.data_preparation.models.pipeline import Pipeline
+
+        if self._pipeline:
+            return self._pipeline
+
+        self._pipeline = Pipeline.get(
+            self.pipeline_uuid,
+            all_projects=True,
+            check_if_exists=False,
+            repo_path=self.repo_path,
+            use_repo_path=True if self.repo_path else False,
+        )
+
+        return self._pipeline
+
+    @property
+    def block(self):
+        if self._block:
+            return self._block
+
+        self._block = self.pipeline.get_block(self.block_uuid)
+
+        return self._block
+
+    @property
+    def variable_uuids(self):
+        return self.block.output_variables(execution_partition=self.execution_partition)
+
+    def get_outputs(self, **kwargs) -> List[Any]:
+        arr = []
+        for variable_uuid in self.variable_uuids:
+            output = self.pipeline.get_block_variable(
+                self.block_uuid,
+                variable_uuid,
+                partition=self.execution_partition,
+            )
+            arr.append(output)
+
+        return arr

mage_ai/data_preparation/models/block/sql/__init__.py

@@ -337,7 +337,7 @@ def execute_sql_code(
             if should_query:
                 return [
                     loader.load(
-                        f'SELECT * FROM {table_name}',
+                        f'SELECT * FROM {schema}.{table_name}',
                         verbose=False,
                     ),
                 ]

mage_ai/data_preparation/models/block/utils.py

@@ -29,6 +29,7 @@ from mage_ai.data_preparation.models.block.dynamic.utils import (
 from mage_ai.data_preparation.models.block.dynamic.utils import (
     should_reduce_output as should_reduce_output_original,
 )
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.constants import (
     DATAFRAME_ANALYSIS_MAX_COLUMNS,
     BlockType,
@@ -393,6 +394,7 @@ def fetch_input_variables(
     dynamic_block_flags: List[DynamicBlockFlag] = None,
     metadata: Dict = None,
     upstream_block_uuids_override: List[str] = None,
+    current_block=None,
 ) -> Tuple[List, List, List]:
     """
     Fetches the input variables for a block.
@@ -510,6 +512,22 @@ def fetch_input_variables(
             if BlockType.GLOBAL_DATA_PRODUCT == upstream_block.type:
                 global_data_product = upstream_block.get_global_data_product()
                 input_vars[idx] = global_data_product.get_outputs()
+
+                mds = {}
+                variable_uuids = upstream_block.output_variables(
+                    execution_partition=execution_partition,
+                )
+                for variable_uuid in variable_uuids:
+                    md = pipeline.get_block_variable(
+                        upstream_block_uuid,
+                        variable_uuid,
+                        partition=execution_partition,
+                    )
+                    if isinstance(md, dict):
+                        mds.update(md)
+
+                kwargs_vars.append(mds)
+
                 continue
 
             # Block output variables for upstream_block_uuid
@@ -743,6 +761,26 @@ def fetch_input_variables(
         __uuid='output_variables'
     )
 
+    if kwargs_vars:
+        kwargs_vars2 = []
+
+        remote_blocks_output = []
+        for kwargs in kwargs_vars:
+            for remote_block_dict in kwargs.get('remote_blocks', []):
+                # Global data products only need the remote block information, not the output
+                if current_block and BlockType.GLOBAL_DATA_PRODUCT == current_block.type:
+                    output = remote_block_dict
+                else:
+                    output = RemoteBlock.load(**remote_block_dict).get_outputs()
+            remote_blocks_output.append(output)
+
+        for kwargs in kwargs_vars:
+            if kwargs.get('remote_blocks'):
+                kwargs['remote_blocks'] = remote_blocks_output
+            kwargs_vars2.append(kwargs)
+
+        kwargs_vars = kwargs_vars2
+
     return input_vars, kwargs_vars, upstream_block_uuids_final
 
 

mage_ai/data_preparation/models/constants.py

@@ -24,6 +24,8 @@ PREFERENCES_FILE = '.preferences.yaml'
 REPO_CONFIG_FILE = 'metadata.yaml'
 VARIABLE_DIR = '.variables'
 
+PIPELINE_RUN_STATUS_LAST_RUN_FAILED = 'last_run_failed'
+
 
 class AIMode(str, Enum):
     OPEN_AI = 'open_ai'

mage_ai/data_preparation/models/global_data_product/__init__.py

@@ -75,6 +75,18 @@ class GlobalDataProduct:
 
         return self._object
 
+    def get_blocks(self) -> List:
+        arr = []
+
+        if not self.settings or len(self.settings) == 0:
+            return arr
+
+        if GlobalDataProductObjectType.PIPELINE == self.object_type:
+            for block_uuid in self.settings.keys():
+                arr.append(self.pipeline.get_block(block_uuid))
+
+        return arr
+
     def get_outputs(self, from_notebook: bool = False, global_vars: Dict = None, **kwargs) -> Dict:
         data = {}
 

mage_ai/data_preparation/models/pipeline.py

@@ -7,7 +7,7 @@ import tempfile
 import zipfile
 from datetime import datetime, timezone
 from io import BytesIO
-from typing import Any, Callable, Dict, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import aiofiles
 import pytz
@@ -62,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
@@ -97,6 +99,7 @@ class Pipeline:
         self.description = description
         self.executor_config = dict()
         self.executor_type = None
+        self._rendered_executor_type = None  # Render template variables
         self.extensions = {}
         self.name = None
         self.notification_config = dict()
@@ -339,7 +342,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')
@@ -514,7 +517,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'):
@@ -673,6 +676,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,
@@ -689,6 +693,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 = []
@@ -715,7 +720,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
 
@@ -726,6 +731,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):
@@ -1294,7 +1309,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())
@@ -1752,10 +1769,13 @@ class Pipeline:
     def get_executable_blocks(self):
         return [b for b in self.blocks_by_uuid.values() if b.executable]
 
-    def get_executor_type(self) -> str:
-        if self.executor_type:
-            return Template(self.executor_type).render(**get_template_vars())
-        return self.executor_type
+    def get_executor_type(self) -> Optional[str]:
+        if self._rendered_executor_type is None:
+            if self.executor_type:
+                return Template(self.executor_type).render(**get_template_vars())
+            else:
+                self._rendered_executor_type = self.executor_type
+        return self._rendered_executor_type
 
     def has_block(self, block_uuid: str, block_type: str = None, extension_uuid: str = None):
         if extension_uuid:
@@ -2132,7 +2152,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)
 
@@ -2205,7 +2225,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/triggers/__init__.py

@@ -206,6 +206,7 @@ def add_or_update_trigger_for_pipeline_and_persist(
     trigger: Trigger,
     pipeline_uuid: str,
     update_only_if_exists: bool = False,
+    old_trigger_name: str = None,
 ) -> Dict:
     trigger_configs_by_name = get_trigger_configs_by_name(pipeline_uuid)
 
@@ -214,13 +215,14 @@ def add_or_update_trigger_for_pipeline_and_persist(
     have, so we need to set "envs" on the updated trigger if it already exists.
     Otherwise, it will get overwritten when updating the trigger in code.
     """
-    existing_trigger = trigger_configs_by_name.get(trigger.name)
+    trigger_name = trigger.name if old_trigger_name is None else old_trigger_name
+    existing_trigger = trigger_configs_by_name.get(trigger_name)
     if existing_trigger is not None:
         trigger.envs = existing_trigger.get('envs', [])
     elif update_only_if_exists:
         return None
 
-    trigger_configs_by_name[trigger.name] = trigger.to_dict()
+    trigger_configs_by_name[trigger_name] = trigger.to_dict()
     yaml_config = dict(triggers=list(trigger_configs_by_name.values()))
     content = yaml.safe_dump(yaml_config)
     trigger_file_path = get_triggers_file_path(pipeline_uuid)

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/local_storage.py

@@ -92,12 +92,18 @@ class LocalStorage(BaseStorage):
             os.makedirs(dirname, exist_ok=True)
 
         with open(file_path, 'w') as file:
-            simplejson.dump(
-                data,
-                file,
-                default=encode_complex,
-                ignore_nan=True,
-            )
+            try:
+                simplejson.dump(
+                    data,
+                    file,
+                    default=encode_complex,
+                    ignore_nan=True,
+                )
+            except ValueError as err:
+                if is_debug():
+                    raise err
+                else:
+                    print(f'[ERROR] LocalStorage.write_json_file: {err}')
 
     async def write_json_file_async(self, file_path: str, data) -> None:
         async with aiofiles.open(file_path, mode='w') as file:

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